TorchScript转Core ML异常：Seq2Seq模型decoder_input_ids未被识别为输入

错误原因分析 #

1. 指定输出shape报错的原因 #

PyTorch转Core ML时，Core ML转换器会根据模型运算逻辑和输入shape自动推导输出shape，手动指定输出shape会与自动推导结果冲突，因此抛出ValueError提示禁止手动指定该参数。

2. decoder_input_ids未被识别为输入的原因 #

Seq2Seq模型包含encoder和decoder双分支，转TorchScript时如果没显式将decoder_input_ids作为前向传播的输入参数，或者TorchScript追踪/脚本化时未捕获到这个输入依赖，Core ML转换器就会忽略该输入。常见场景：

• 模型前向传播中，decoder_input_ids是内部生成（比如用默认值或其他输入推导），而非外部传入
• TorchScript追踪时的示例输入未包含decoder_input_ids，导致转换器判定该参数不是必需输入

解决办法 #

步骤1：修正TorchScript转换，确保捕获decoder_input_ids输入 #

必须让模型前向函数显式接收decoder_input_ids作为参数，且追踪时传入对应示例输入：

import torch
from transformers import T5ForConditionalGeneration

# 加载预训练Seq2Seq模型（以T5为例）
model = T5ForConditionalGeneration.from_pretrained("t5-small")
model.eval()

# 包装模型，显式声明所有输入参数
class WrappedModel(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask):
        return self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask
        ).logits

wrapped_model = WrappedModel(model)

# 准备包含所有输入的示例张量
batch_size = 1
seq_len = 16
decoder_seq_len = 16
input_ids = torch.randint(0, model.config.vocab_size, (batch_size, seq_len))
attention_mask = torch.ones((batch_size, seq_len))
decoder_input_ids = torch.randint(0, model.config.vocab_size, (batch_size, decoder_seq_len))
decoder_attention_mask = torch.ones((batch_size, decoder_seq_len))

# 追踪生成TorchScript模型
traced_model = torch.jit.trace(wrapped_model, (input_ids, attention_mask, decoder_input_ids, decoder_attention_mask))
traced_model.save("traced_seq2seq_model.pt")

步骤2：转换为Core ML，依赖自动推导输出shape #

使用coremltools转换时，直接传入TorchScript模型和输入描述，无需指定输出shape：

import coremltools as ct

# 加载TorchScript模型
traced_model = torch.jit.load("traced_seq2seq_model.pt")

# 定义输入类型与名称
input_descriptions = [
    ct.TensorType(name="input_ids", shape=input_ids.shape, dtype=ct.int32),
    ct.TensorType(name="attention_mask", shape=attention_mask.shape, dtype=ct.int32),
    ct.TensorType(name="decoder_input_ids", shape=decoder_input_ids.shape, dtype=ct.int32),
    ct.TensorType(name="decoder_attention_mask", shape=decoder_attention_mask.shape, dtype=ct.int32)
]

# 转换为Core ML模型（推荐用MLProgram格式，支持iOS 15+）
coreml_model = ct.convert(
    traced_model,
    inputs=input_descriptions,
    convert_to="mlprogram"
)

coreml_model.save("seq2seq_coreml.mlmodel")

额外验证与修复 #

用Xcode打开生成的.mlmodel文件，检查输入列表是否包含decoder_input_ids。若仍有问题：

• 改用torch.jit.script替代追踪，确保捕获所有动态输入逻辑
• 检查模型前向传播的条件分支，确保decoder_input_ids在所有路径中都被使用，避免被优化器忽略

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