PyTorch模型导出ONNX必须满足可追踪性、显式设model.eval()、正确配置input_names/output_names和dynamic_axes、选用兼容opset_version,并用ONNX Runtime验证数值一致性。

PyTorch模型能直接导出为ONNX,但必须满足“可追踪”条件——即模型前向传播不能含动态控制流、不可导操作或Python内置类型返回值;否则torch.onnx.export会静默截断或报错。
导出前必须确保模型处于eval模式且无训练专用层
Dropout、BatchNorm在train模式下行为与推理不一致,ONNX Runtime运行时会因输入维度/统计量不匹配而崩溃。即使模型本身没调用train(),也要显式设为eval():
-
model.eval()是硬性要求,不是建议 - 若模型含
nn.Dropout但未设eval(),导出的ONNX中该层仍保留随机丢弃逻辑,部署后输出不稳定 - 自定义模块中若用了
self.training做分支判断(如某些GAN判别器),需临时重写forward或用torch.no_grad()包裹,否则ONNX无法捕获完整路径
torch.onnx.export的关键参数不能只靠默认值
最常踩的坑是忽略input_names/output_names和dynamic_axes,导致后续推理时张量绑定失败:
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac,无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启:辅助功能 + 屏幕录制权限
-
input_names=["input"]和output_names=["output"]必须显式指定,否则ONNX文件内部节点名是自动生成的(如123、456),下游框架无法按名取输出 - 若输入含batch维度(如
[B, 3, 224, 224]),必须用dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}声明,否则ONNX Runtime加载后固定为B=1,换其他batch size直接报错 -
opset_version建议选14或15:太低(如11)不支持torch.where等常用算子;太高(如17)部分旧版ONNX Runtime不兼容
导出后必须用ONNX Runtime验证shape和数值一致性
导出成功不等于可用。常见现象是ONNX输出shape对得上但数值偏差超1e-4,根源多在PyTorch和ONNX对算子实现的细微差异:
立即学习“Python免费学习笔记(深入)”;
- 用
onnxruntime.InferenceSession加载导出的.onnx文件,传入和PyTorch前向完全相同的numpy输入(注意dtype:PyTorch float32 → numpy.float32,别用float64) - 对比PyTorch输出
y_pt = model(x)和ONNX输出y_onnx = session.run(None, {"input": x.numpy()})[0],用np.allclose(y_pt.detach().numpy(), y_onnx, atol=1e-5)验证 - 若不通过,优先检查模型中是否用了
torch.tensor([1,2,3])这类非参数张量——它们在ONNX里可能被常量化掉,导致数值偏移;应改用nn.Parameter或注册为buffer
真正麻烦的是那些不报错但结果漂移的case:比如torch.nn.functional.interpolate在mode="bilinear" + align_corners=False时,PyTorch和ONNX的插值核实现有微小差异,必须在导出时加keep_initializers_as_inputs=True并手动校准。

















