3、模型导出实战：从PyTorch导出ONNX模型（torch.onnx.export），常见参数详解
好，咱们直接进入正题。今天聊的是PyTorch模型转ONNX最核心的一步——torch.onnx.export。这个函数，说白了就是PyTorch和ONNX之间的桥梁。我刚开始接触ONNX时，觉得这玩意儿不就是一行代码的事吗？后来发现，参数调不好，导出的模型要么跑不起来，要么精度对不上。嗯，这里面的门道不少。
3.1 最基础的导出流程
先看一个最简单的例子。假设你有一个训练好的PyTorch模型，想导出成ONNX：
import torch
import torch.onnx
# 假设 model 是你的 PyTorch 模型
model = MyModel()
model.eval() # 记得切到 eval 模式
# 构造一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)
# 导出
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 模型的输入张量
"model.onnx", # 输出文件名
input_names=['input'],
output_names=['output']
)
这段代码跑完，你就能得到一个 model.onnx 文件。但说实话，这只是最基础的用法。实际项目中，你肯定需要调更多参数。
model.eval()。为什么？因为 dropout 和 batch normalization 在训练和推理时的行为不一样。我见过有人忘了这步，结果导出的 ONNX 推理结果和 PyTorch 对不上，排查了半天才发现是这个问题。
3.2 核心参数详解
咱们一个一个来看。这些参数，我按重要程度排了个序。
3.2.1 opset_version
这个参数指定了 ONNX 算子集的版本。说白了，就是告诉导出器：你想用哪个版本的 ONNX 标准。
torch.onnx.export(
model, dummy_input, "model.onnx",
opset_version=11 # 常用版本
)
我个人建议用 opset_version=11 或 13。为什么？
- opset 11 支持了大多数常用算子，兼容性很好
- opset 13 增加了对
Softmax等算子的优化 - opset 17 虽然新，但有些推理引擎还没完全支持
3.2.2 input_names 和 output_names
这两个参数给模型的输入输出张量起名字。别小看这个，后面做推理时，你要靠这些名字来获取数据。
torch.onnx.export(
model, dummy_input, "model.onnx",
input_names=['input'], # 输入张量名字
output_names=['output'] # 输出张量名字
)
你想想看，如果一个模型有多个输入（比如图片和文本），你给它们起名叫 image_input 和 text_input，后面调试时一眼就能看出来哪个是哪个。
3.2.3 dynamic_axes
这个参数太重要了。它用来指定哪些维度是动态的。比如，你的模型输入 batch size 是可变的：
torch.onnx.export(
model, dummy_input, "model.onnx",
dynamic_axes={
'input': {0: 'batch_size'}, # 第0维是动态的
'output': {0: 'batch_size'}
}
)
如果不设置 dynamic_axes，导出的 ONNX 模型会固定输入输出的形状。比如你传了 (1, 3, 224, 224)，那模型就只接受这个形状。想换个 batch size？不行。
3.2.4 export_params
这个参数控制是否导出模型权重。默认是 True。如果你只想导出模型结构，可以设为 False：
torch.onnx.export(
model, dummy_input, "model.onnx",
export_params=False # 只导出结构，不导出权重
)
嗯，这个场景比较少见。我一般只在做模型结构分析时才会用到。
3.2.5 do_constant_folding
常量折叠优化。说白了，就是把一些在推理时不会变化的计算提前算好，存成常量。
torch.onnx.export(
model, dummy_input, "model.onnx",
do_constant_folding=True # 开启常量折叠
)
我建议默认开启。它能减小模型体积，还能提升推理速度。但要注意，如果模型中有一些看似常量但实际上会变化的操作（比如某些自定义层），开启后可能导致结果错误。
3.2.6 verbose
调试利器。设为 True 后，导出过程中会打印详细的算子转换信息：
torch.onnx.export(
model, dummy_input, "model.onnx",
verbose=True # 打印详细信息
)
当导出失败时，这个参数能帮你快速定位是哪个算子出了问题。
3.3 一个完整的导出示例
好了，咱们把上面这些参数串起来，写一个实际项目中常用的导出代码：
import torch
import torch.onnx
def export_to_onnx(model, dummy_input, save_path):
model.eval()
torch.onnx.export(
model,
dummy_input,
save_path,
opset_version=11,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
},
do_constant_folding=True,
export_params=True,
verbose=False
)
print(f"模型已导出至: {save_path}")
# 使用示例
model = MyModel()
dummy = torch.randn(1, 3, 224, 224)
export_to_onnx(model, dummy, "my_model.onnx")
3.4 导出后的验证
导出完别急着走。我建议做两件事：
- 用 ONNX Runtime 跑一遍推理，看结果和 PyTorch 是否一致
- 用 Netron 打开 ONNX 文件，可视化检查模型结构
验证代码大概长这样：
import onnxruntime as ort
import numpy as np
# 加载 ONNX 模型
session = ort.InferenceSession("my_model.onnx")
# 准备输入
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 推理
outputs = session.run(None, {input_name: input_data})
print("ONNX 推理结果:", outputs[0])
3.5 常见问题与解决方案
| 问题 | 原因 | 解决方案 |
|---|---|---|
| 导出时报 "Unsupported operator" | PyTorch 的某个算子 ONNX 不支持 | 升级 opset 版本，或自定义算子映射 |
| ONNX 推理结果和 PyTorch 不一致 | 模型未切到 eval 模式，或动态轴设置错误 | 检查 model.eval()，确认 dynamic_axes 配置 |
| 导出的 ONNX 文件太大 | 模型权重未压缩，或常量折叠未开启 | 开启 do_constant_folding，或使用 ONNX 量化 |
| 推理时输入形状不匹配 | 未设置 dynamic_axes，或设置错误 | 检查 dynamic_axes 配置，确保维度索引正确 |
嗯，以上就是 torch.onnx.export 的核心内容。说白了，导出模型就像打包行李——参数调好了，一路顺畅；参数没调好，后面全是坑。我建议你每次导出后都做一遍验证，养成习惯。
下一章，咱们聊聊如何用 ONNX Runtime 做推理，以及怎么排查推理中的精度问题。