torch.onnx.export导出Tacotron模型报ValueError解包值过多错误

报错根因 #

模型forward方法约定输入为包含5个张量的可迭代结构，分别对应text_inputs、text_lengths、mels、max_len、output_lengths五个入参。当前代码传入的dummy_input是单个5维张量，PyTorch执行解包逻辑时会沿张量第0维拆分，最终得到30个拆分元素，和预期的5个入参数量不匹配，触发值错误。
另外现有导出代码还存在两个隐性问题：

• 导出前未将模型切换到评估模式，训练态的dropout、BatchNorm层会导致导出模型推理结果异常
• input_names、output_names参数传入了单个字符串，不符合接口要求的字符串列表格式，修复解包问题后会触发新的参数报错

修复步骤 #

1. 构造符合模型输入契约的dummy输入
   按照模型实际推理时的输入形状、数据类型、设备，构造5个独立张量，封装为元组作为导出输入，示例代码如下（张量形状请根据自身模型的实际配置调整，不要直接硬抄）：

   model = model.cuda().eval()
   batch_size = 1
   # 替换为你模型实际使用的输入维度、dtype
   text_inputs = torch.randint(0, 100, (batch_size, 200), dtype=torch.long).cuda()
   text_lengths = torch.tensor([200], dtype=torch.long).cuda()
   mels = torch.randn(batch_size, 80, 500, dtype=torch.float32).cuda()
   max_len = torch.tensor(1000, dtype=torch.long).cuda()
   output_lengths = torch.tensor([500], dtype=torch.long).cuda()
   dummy_input = (text_inputs, text_lengths, mels, max_len, output_lengths)
   

2. 修正导出接口参数
   调整input_names、output_names为字符串列表格式，指定兼容后续TFLite转换的opset版本，开启常量折叠优化：

   torch.onnx.export(
       model,
       dummy_input,
       "myModel.onnx",
       verbose=False,
       input_names=["text_inputs", "text_lengths", "mels", "max_len", "output_lengths"],
       output_names=["mel_pred", "stop_token_pred", "alignment"], # 按实际输出数量、命名调整
       opset_version=16,
       do_constant_folding=True,
   )
   

后续转TFLite注意事项 #

• 导出ONNX后先通过ONNX Runtime做一次推理校验，确认输出和PyTorch原生推理结果误差在可接受范围内，再进行TFLite转换
• 若模型包含自回归循环解码逻辑，建议将编码器、解码器拆分为两个独立ONNX模型分别导出，避免控制流算子导致TFLite转换失败
• opset版本选择13~16区间的版本兼容性最佳，过高或过低都可能出现算子不支持的问题

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