含位置与关键字参数的PyTorch Module导出ONNX参数错误排查

错误原因 #

PyTorch 1.8.0中torch.onnx.export的args参数会将其中每个元素都作为位置参数传递给模块的forward方法。你传入args=(cond张量, {'input': x})时，相当于直接调用model(cond张量, {'input': x})——此时forward会收到3个位置参数（隐含的self、cond张量、字典{'input':x}），但你的forward定义只接受2个显式位置参数（self和cond），剩余参数需要通过**kwargs接收关键字参数，因此触发参数数量不匹配的错误。

正确传参方式 #

有两种可行的修正方案：

方案1：使用onnx.export的kwargs参数传递关键字参数

保持forward定义不变，在导出时将位置参数放在args中，关键字参数单独传入kwargs参数：

torch.onnx.export(
    model,
    args=(cond张量,),  # 仅传递位置参数cond
    kwargs={'input': x},  # 传递forward需要的关键字参数
    # 其他导出参数（如output_path、opset_version等）
)

这样onnx.export会正确调用model(cond张量, input=x)，匹配forward(self, cond, **kwargs)的参数定义。

方案2：调整args为直接模拟正常调用的参数形式

如果不想使用kwargs参数，也可以将关键字参数对应的张量直接作为位置参数传入args，但需要同步修改forward的定义，把**kwargs改为具体的关键字参数：

# 修改forward定义
def forward(self, cond, input):
    # 原有逻辑
    ...

# 导出时传入对应位置参数
torch.onnx.export(
    model,
    args=(cond张量, x),
    # 其他导出参数
)

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