1、PyTorch模型导出基础：理解torch.onnx.export的核心参数与导出流程
说实话，模型部署这件事，我踩过的坑比走过的路还多。刚开始做PyTorch转ONNX时，总觉得不就是一行代码的事吗？后来才发现，这里面的门道真不少。今天咱们就从头捋一捋，把torch.onnx.export这个核心函数彻底搞明白。
1.1 为什么需要ONNX？
你想想看，PyTorch模型训练好了，总不能一直活在Python环境里吧？生产环境可能是C++、Java，甚至跑在手机或嵌入式设备上。ONNX就是那个「通用翻译官」——它定义了一套标准化的模型表示格式，让不同框架之间能互相转换。
我个人习惯把ONNX比作「模型界的PDF」。你写Word文档，别人不一定有Word，但PDF几乎谁都能打开。ONNX也是这个道理，PyTorch训练完，导出成ONNX，然后TensorRT、OpenVINO、ONNX Runtime这些推理引擎都能直接吃进去。
核心价值：一次导出，到处推理。跨平台、跨语言、跨硬件。
1.2 torch.onnx.export 的核心参数
这个函数我用了不下几百次了。每次新项目，第一件事就是调它。咱们直接看最常用的参数：
import torch
import torch.onnx
# 假设我们有一个训练好的模型
model = torch.load('my_model.pth')
model.eval()
# 构造一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)
# 导出
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 示例输入
'model.onnx', # 输出文件名
export_params=True, # 是否导出模型参数
opset_version=11, # ONNX算子集版本
do_constant_folding=True, # 是否做常量折叠优化
input_names=['input'], # 输入张量名称
output_names=['output'], # 输出张量名称
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
嗯，这里要注意几个关键点：
| 参数 | 作用 | 我的建议 |
|---|---|---|
model |
要导出的模型实例 | 记得先调 model.eval()，去掉Dropout和BatchNorm的训练行为 |
dummy_input |
一个形状正确的示例输入 | 形状要和实际推理时一致，否则导出后可能报错 |
opset_version |
ONNX算子集版本号 | 我一般用11或13，太老不支持新算子，太新有些推理引擎不兼容 |
dynamic_axes |
指定哪些维度是动态的 | 如果推理时batch size会变，这个必须配 |
1.3 导出流程：从模型到ONNX
整个流程其实就三步，但每一步都有坑。我当年第一次导出时，卡在第二步整整两天。
- 准备模型和输入：加载模型，调成eval模式，构造一个dummy input
- 调用export函数：传入模型、输入、输出路径和参数
- 验证ONNX模型：用
onnx.checker检查，再用ONNX Runtime跑一遍看看结果对不对
小技巧：导出后一定要做验证。我曾经遇到过导出成功但推理结果完全不对的情况，就是因为某个自定义算子没注册好。
1.4 避坑指南：我踩过的那些坑
我曾经在一个图像分割项目里，模型导出后死活跑不通。查了两天，最后发现是torch.nn.functional.interpolate这个函数在ONNX里不支持某些参数组合。后来我换成了torch.nn.Upsample层，问题就解决了。
还有一次，我忘了调model.eval()，结果导出的ONNX模型里包含了Dropout层，推理时每次结果都不一样。你说这坑不坑？
常见坑点：
- 忘记调
model.eval()—— 导出结果包含训练行为 - dummy_input 形状和实际不一致 —— 导出后推理报错
- 使用了ONNX不支持的算子 —— 需要手动替换或自定义
- 动态轴没配好 —— 推理时batch size固定死
1.5 验证导出的ONNX模型
导出完别急着走，一定要验证。我习惯写个脚本：
import onnx
import onnxruntime as ort
import numpy as np
# 1. 检查ONNX模型结构
onnx_model = onnx.load('model.onnx')
onnx.checker.check_model(onnx_model)
print('ONNX模型结构检查通过')
# 2. 用ONNX Runtime推理
ort_session = ort.InferenceSession('model.onnx')
ort_inputs = {ort_session.get_inputs()[0].name: np.random.randn(1, 3, 224, 224).astype(np.float32)}
ort_outputs = ort_session.run(None, ort_inputs)
print('ONNX Runtime推理成功，输出形状:', ort_outputs[0].shape)
说白了，这一步就是给导出的模型做个「体检」。结构没问题，推理能跑通，这才算真正导出成功。
1.6 总结一下
PyTorch转ONNX，核心就是torch.onnx.export这个函数。你只要记住：
- 模型要调
eval() - dummy input形状要对
- opset版本选合适
- 动态轴按需配置
- 导出后一定要验证
做到这几点，基本就不会出大问题。下一章咱们聊聊更复杂的场景——带自定义算子的模型怎么导出。那个才是真正考验功力的时候。