Seq2Seq模型导出ONNX后解码器输入长度适配问题求助

• 排查模型内硬编码的序列长度
  报错显示输入形状{2,1,300}无法reshape为{1,20,15}，说明模型内部存在硬编码的固定长度（比如20）或batch_size（比如1）。检查所有涉及reshape、线性层维度计算的代码：

  • 替换所有固定数值为动态获取的维度，比如用input_seq.size(1)获取当前输入序列长度，用input_seq.size(0)获取batch_size。
  • 示例修正：将x.reshape(1, 20, 15)改为x.reshape(input_seq.size(0), input_seq.size(1), hidden_dim)（需确保hidden_dim与模型隐藏层维度匹配，保证维度计算的一致性）。

• 确保动态轴设置覆盖所有相关维度
  即使设置了动态轴，若漏设关键维度仍会导致问题。导出解码器时，需明确标记所有动态变化的维度：

  torch.onnx.export(
      decoder,
      (decoder_input, encoder_output, attn_mask),
      "decoder.onnx",
      dynamic_axes={
          "decoder_input": {0: "batch_size", 1: "tgt_seq_len"},
          "encoder_output": {0: "batch_size", 1: "src_seq_len"},
          "attn_mask": {1: "tgt_seq_len", 2: "src_seq_len"},
          "decoder_output": {0: "batch_size", 1: "tgt_seq_len"}
      },
      opset_version=17,  # 建议使用较高版本opset以支持更多动态操作
      verbose=False
  )
  

  确认所有输入、输出以及中间依赖的张量（如注意力掩码）的可变维度都被标记为动态。

• 使用TorchScript而非Trace导出模型
  若之前用torch.jit.trace导出模型，它会记录导出时的固定形状，无法保留动态逻辑。改用torch.jit.script先将模型转换为脚本模块，再导出ONNX：

  # 先将解码器转为ScriptModule
  scripted_decoder = torch.jit.script(decoder)
  # 再导出ONNX
  torch.onnx.export(
      scripted_decoder,
      (decoder_input, encoder_output, attn_mask),
      "decoder.onnx",
      dynamic_axes=...,  # 沿用上述动态轴设置
      opset_version=17
  )
  

  TorchScript能正确保留模型中的动态分支、循环和维度计算逻辑，避免Trace带来的固定形状问题。

• 检查注意力掩码的动态生成逻辑
  若之前固定了解码器输入长度并手动生成掩码，需改为在forward函数内根据当前输入序列长度动态生成掩码：

  def forward(self, tgt, enc_output, src_mask=None):
      tgt_seq_len = tgt.size(1)
      # 动态生成自注意力掩码
      tgt_mask = self.generate_square_subsequent_mask(tgt_seq_len).to(tgt.device)
      # 后续注意力计算使用动态生成的掩码
      ...
  

  确保掩码的尺寸与当前输入的序列长度完全匹配，避免因掩码尺寸固定导致的注意力计算错误。

• 验证ONNX模型的合法性
  用ONNX工具检查模型是否存在错误：

  import onnx
  model = onnx.load("decoder.onnx")
  onnx.checker.check_model(model)  # 检查模型结构合法性
  

  也可以用Netron可视化模型，查看是否存在固定维度的节点，确认动态轴是否被正确应用。

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