Python怎么把PyTorch模型导出为ONNX格式_torch.onnx.export与dynamic_axes设定

Python怎么把PyTorch模型导出为ONNX格式_torch.onnx.export与dynamic_axes设定

导出时模型不报错但推理结果对不上，大概率是 dynamic_axes 没设对

很多开发者都踩过这个坑：PyTorch模型导出ONNX时一切顺利，但一到推理环节，结果要么对不上，要么直接报错。问题根源，十有八九出在dynamic_axes这个参数上。

要知道，PyTorch导出ONNX时，默认会把所有张量的形状当作固定值来处理。一旦你的模型涉及动态批次大小、变长序列（比如NLP中处理不同长度的句子），或者内部包含条件分支结构（例如if x.size(0) > 1:），如果没有正确声明dynamic_axes，ONNX Runtime在推理时就会按照固定的形状去处理，导致形状不匹配，甚至产生难以察觉的数值偏差。这种情况在使用torch.nn.functional.pad、torch.where或自定义forward逻辑时尤为常见。

所以，关键不在于“要不要设置”，而在于“哪些维度必须标记”。以下是几个必须关注的常见动态维度：

• 输入的第0维（批次大小）：这几乎是标配，通常需要标记为{'input': {0: 'batch'}}。
• NLP模型的第1维（序列长度）：例如{'input_ids': {1: 'seq_len'}}。
• 检测模型中的边界框数量维度（通常是第1维）：比如{'boxes': {1: 'num_boxes'}}。
• 输出张量：如果其形状依赖于输入（例如Mask R-CNN中的掩码输出），对应的动态维度也必须同步标记。

torch.onnx.export 的 input_names / output_names 和 dynamic_axes 必须严格一致

这是一个非常严格的对应关系，但很容易被忽略。input_names和output_names定义了ONNX计算图中张量的符号名称，而dynamic_axes字典的键（key）必须与这些名称完全一致。任何细微的差别——比如大小写错误、多了一个空格，或者误用了模型内部的变量名（如'x'）而不是你传入的参数名（如'input'）——都会导致dynamic_axes设置失效。最终导出的ONNX模型，其输入输出形状依然显示为固定值。

具体怎么操作？记住这几个要点：

• 当使用input_names=['input']时，dynamic_axes只能写成{'input': {0: 'batch'}}，而不能是{'x': ...}。
• 对于多个输入，需要为每个输入名称单独列出动态轴。例如：input_names=['input_ids', 'attention_mask'] 对应 dynamic_axes={'input_ids': {1: 'seq_len'}, 'attention_mask': {1: 'seq_len'}}。
• 如果不确定输出的名称，一个实用技巧是：先用torch.onnx.export(..., verbose=True)导出一次，查看日志中打印出的输出名称，然后直接复制过来使用。

导出失败报错 “Unsupported value type” 或 “Cannot export function”

这类错误信息通常指向一个不同的问题：模型中使用了ONNX标准不支持的运算符或操作。虽然它和dynamic_axes无关，但很容易被混淆。以下几个是典型的触发点：

• 运行时设备指定：例如torch.tensor([1, 2, 3], device='cuda')。ONNX不支持在运行时指定设备。应改为torch.tensor([1, 2, 3]).to(model.device)或直接使用CPU张量。
• 张量形状作为参数：在torch.arange或torch.linspace中使用了非整数/浮点数标量，例如torch.arange(x.shape[0])中的x.shape[0]是torch.Size类型。应改为torch.arange(x.size(0))。
• Python原生控制流：在自定义的__call__或重载的forward方法中，使用了基于张量值的Pythonfor循环或if判断。需要替换为torch.where、torch.masked_fill等可以被导出的操作符。
• 模型模式：使用了torch.jit.script包装但未提前设置model.eval()。务必在导出ONNX前调用model.eval()，并禁用dropout和批归一化层的更新。

验证 dynamic_axes 是否生效：别只看 export 成功，要看 onnxruntime 推理时能否接受不同尺寸输入

导出过程没有报错，绝不意味着dynamic_axes已经正确生效。最直接、最可靠的验证方法是：使用ONNX Runtime加载模型后，主动喂入不同形状的输入数据进行测试。

import onnxruntime as ort

sess = ort.InferenceSession('model.onnx')

# 先测试 batch_size=1
out1 = sess.run(None, {'input': np.random.randn(1, 3, 224, 224).astype(np.float32)})

# 再测试 batch_size=4 —— 如果这一步报错 “Input shape mismatch”，就说明 dynamic_axes 没有起作用
out4 = sess.run(None, {'input': np.random.randn(4, 3, 224, 224).astype(np.float32)})

如果第二步测试失败了，请立刻回头检查：dynamic_axes的键名拼写是否与input_names完全一致、是否遗漏了某些输入、导出时提供的example_inputs形状是否正确（例如应该是torch.randn(1, 3, 224, 224)而不是torch.randn(4, 3, 224, 224)）。

更棘手的是那些“看似生效但数值漂移”的情况——比如序列长度变化后，Softmax层的输出概率与PyTorch原始结果不一致。这类问题往往藏得比较深，可能需要逐层比对中间张量的输出。常见的罪魁祸首包括padding mask的广播逻辑，或者position embedding的索引计算方式在动态形状下出现了偏差。