使用torch.jit.trace追踪PyTorch模型时遇RuntimeError问题求助

尝试将PyTorch模型转换为ScriptModule以便部署到C++环境，执行torch.jit.trace时触发上述错误，相关代码见下文。

错误原因 #

1. trace调用方式错误：直接传入了model_rnn.forward(...)的执行结果（张量），而非模型实例或可追踪的可调用对象。torch.jit.trace需要接收模型/函数作为第一个参数，再传入示例输入来追踪执行流程。
2. forward参数设计问题：forward函数将device作为参数传入，这类非张量类型输入在torchscript中支持有限，容易引发追踪失败；同时seq_len为Python列表，也可能导致追踪异常。

解决方案 #

步骤1：修改模型结构，优化forward参数 #

移除forward和encoder中的device参数，直接从输入张量获取设备；调整参数传递逻辑：

from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence
from .. import TORCH_DEVICE, __version_ft__, __version_train__

class LmkRNN(nn.Module):
    def __init__(self, feature_size, hidden_size, out_size, n_layer=2, bf=False, bi=False):
        super(LmkRNN, self).__init__()
        self.hidden_size = hidden_size
        self.n_layer = n_layer
        self.batch_size = 0
        self.max_seq_len = 0
        self.size_bl_out = feature_size * 2
        self.n_direction = 1
        self.rnn_en = nn.GRU(input_size=self.size_bl_out, hidden_size=hidden_size,
                             num_layers=n_layer, batch_first=bf, bidirectional=bi)
        self.h_en = 0
        self.fc1 = nn.Linear(hidden_size, out_size)

    def encoder(self, seq, seq_len):
        max_seq_len, batch_size, n_pts, _ = seq.shape
        self.batch_size = batch_size
        self.max_seq_len = max_seq_len
        seq_c = torch.tanh(torch.flatten(seq, start_dim=-2))
        packed = pack_padded_sequence(
            seq_c, seq_len, batch_first=False, enforce_sorted=False)
        # 从输入张量获取设备，无需显式传入
        device = seq.device
        self.h_en = torch.zeros(
            (self.n_layer, self.batch_size, self.hidden_size), device=device)
        out, h_n = self.rnn_en(packed, self.h_en)
        self.h_en = h_n

    def forward(self, seq, seq_len):
        self.encoder(seq=seq, seq_len=seq_len)
        enc = self.h_en[-1]
        out = self.fc1(enc)
        return F.normalize(out)

步骤2：修正torch.jit.trace调用方式 #

将seq_len转为张量，传入模型实例和匹配的示例输入：

# 模型参数定义（保持不变）
n_lmk_pts = 476
hidden_size = 32
out_size = 64
n_layers = 2
FACE_EMB_SIZE = 512
device = torch.device(TORCH_DEVICE)
model_rnn = LmkRNN(feature_size=n_lmk_pts, hidden_size=hidden_size, out_size=out_size, n_layer=n_layers).to(device).float()

def compute_lmkemb(lmk_seq):
    print("In compute_lmkemb")
    with torch.no_grad():
        # 将seq_len转为张量，适配torchscript
        seq_len = torch.tensor([lmk_seq.shape[0], ], dtype=torch.int64).to(device)
        tensors = [torch.from_numpy(lmk_seq), ]
        ps = pad_sequence(tensors, batch_first=False)
        lmkdata = ps.float().to(device)
        emb = model_rnn(lmkdata, seq_len)

        # 正确调用trace：传入模型实例+匹配的示例输入元组
        traced_script_module = torch.jit.trace(model_rnn, example_inputs=(lmkdata, seq_len))

        # 保存ScriptModule
        filename = "scriptmodule.pt"
        traced_script_module.save(filename)
        print(f"Successfully created scriptmodule file {filename}.")

    out = emb[0].detach().cpu().numpy()
    return out

补充说明 #

• 若需要保留seq_len为列表，trace时也可直接传入列表作为示例输入，但转为张量更符合torchscript的规范，能避免潜在问题。
• 若模型有预训练权重，需确保trace前已完成权重加载，trace会记录当前模型的权重状态。

内容的提问来源于stack exchange，提问作者t_perk