PyTorch导出ONNX及dynamic_axes设置详解

PyTorch导出ONNX模型时，推理结果与原模型不一致的“隐形杀手”往往不是算子兼容性问题，而是dynamic_axes设置不当——ONNX默认将所有张量shape视为固定值，若未显式声明batch、序列长度、检测框数量等动态维度，ONNX Runtime在处理变长输入时就会出现shape不匹配或数值偏差；本文直击痛点，详解如何精准标记input_names/output_names中对应维度（如{'input': {0: 'batch'}, 'input_ids': {1: 'seq_len'}}），强调命名一致性、常见易错场景（如device指定、Python控制流、arange使用）及关键验证方法（用不同尺寸输入实测ONNX Runtime响应），帮你避开导出成功却推理翻车的深坑。

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

PyTorch 导出 ONNX 时默认把所有 tensor shape 当作固定值处理。一旦模型里有动态 batch、变长序列（比如 NLP 中不同长度的句子）、或带条件分支的结构（如 if x.size(0) > 1:），没声明 dynamic_axes 就会导致 ONNX runtime 推理时 shape 不匹配，或者数值偏差——尤其在使用 torch.nn.functional.pad、torch.where 或自定义 forward 逻辑时。

关键不是“要不要设”，而是“哪些维度必须标”。常见需标记的维度：

• input 的第 0 维（batch size）：几乎所有情况都要标 {'input': {0: 'batch'}}
• NLP 模型的第 1 维（sequence length）：如 {'input_ids': {1: 'seq_len'}}
• 检测模型中 bbox 数量维度（常为第 1 维）：如 {'boxes': {1: 'num_boxes'}}
• 输出 tensor 若 shape 依赖输入（如 maskrcnn 的 mask 输出），对应维度也得同步标

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

input_names 和 output_names 是 ONNX 图里 tensor 的符号名，dynamic_axes 的 key 必须完全等于其中某个名字。写错大小写、多空格、或用内部变量名（如 'x'）代替传入的参数名（如 'input'），都会让 dynamic_axes 失效——ONNX 生成后仍显示为固定 shape。

实操建议：

• 用 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) 看日志里打印的 output names，再抄过来

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

这类错误通常和模型里用了 ONNX 不支持的操作有关，不是 dynamic_axes 的问题，但容易被误判。典型触发点：

• torch.tensor([1, 2, 3], device='cuda') —— ONNX 不支持运行时 device 指定，改用 torch.tensor([1, 2, 3]).to(model.device) 或直接 CPU tensor
• torch.arange / torch.linspace 带非 int/float 标量（如 torch.arange(x.shape[0]) 中 x.shape[0] 是 torch.Size 类型）→ 改成 torch.arange(x.size(0))
• 自定义 __call__ 或重载 forward 里用了 Python 控制流（for, if 判 tensor 值）→ 改用 torch.where、torch.masked_fill 等可导出操作
• 用了 torch.jit.script 包装但没提前 model.eval() → ONNX 导出前务必调用 model.eval() 并禁用 dropout/bn 更新

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

导出成功不代表 dynamic_axes 生效。最直接的验证方式是用 onnxruntime.InferenceSession 加载模型后，主动 feed 不同 shape 的输入：

import onnxruntime as ort
sess = ort.InferenceSession('model.onnx')
# 先试 batch=1
out1 = sess.run(None, {'input': np.random.randn(1, 3, 224, 224).astype(np.float32)})
# 再试 batch=4 —— 如果这里报错 "Input shape mismatch"，说明 dynamic_axes 没起作用
out4 = sess.run(None, {'input': np.random.randn(4, 3, 224, 224).astype(np.float32)})

如果第二步失败，回头检查：dynamic_axes key 是否拼写一致、是否漏了 input_names、导出时是否传了 example_inputs 对应的正确 shape（例如 torch.randn(1, 3, 224, 224) 而不是 torch.randn(4, 3, 224, 224)）。

真正麻烦的是那些“看起来生效了但数值漂移”的情况——比如 sequence length 变化后 softmax 输出概率和 PyTorch 不一致。这时候要逐层比对中间 tensor，往往卡在 padding mask 的广播逻辑或 position embedding 的索引计算上。

文中关于的知识介绍，希望对你的学习有所帮助！若是受益匪浅，那就动动鼠标收藏这篇《PyTorch导出ONNX及dynamic_axes设置详解》文章吧，也可关注golang学习网公众号了解相关技术文章。