3、PyTorch模型导出ONNX：torch.onnx.export详解、动态轴与静态轴设置、输入输出命名规范、opset版本选择
好，咱们进入第三章。前面两章我们把环境搭好了，也跑通了YOLO的推理。但说实话，PyTorch模型在生产环境里直接跑，效率上不去，部署也麻烦。所以这一章，咱们得把PyTorch模型转成ONNX——一个中间格式，方便后续做TensorRT优化。
torch.onnx.export这个函数，我刚开始用的时候也踩过不少坑。说白了，它就是把PyTorch的计算图“拍”成ONNX的静态图。但怎么拍得对、拍得好，这里头门道不少。
3.1 torch.onnx.export核心参数解析
先看一个最基础的导出代码。我习惯把参数都写清楚，方便后面调试。
import torch
import torch.onnx
# 假设我们有一个训练好的YOLOv8模型
model = torch.load('yolov8n.pt', map_location='cpu')
model.eval()
# 创建一个虚拟输入，batch_size=1，3通道，640x640
dummy_input = torch.randn(1, 3, 640, 640)
# 导出ONNX
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 虚拟输入
'yolov8n.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 |
要导出的PyTorch模型实例 | 务必先调用model.eval()，关掉dropout和batch norm的training模式 |
dummy_input |
一个形状正确的虚拟输入张量 | 用torch.randn生成，数值不重要，形状要对 |
export_params |
是否导出模型权重 | 默认True，一般不动 |
opset_version |
ONNX算子集版本 | YOLO推荐11或12，后面细说 |
do_constant_folding |
是否做常量折叠优化 | 建议True，能减小模型体积 |
input_names |
输入张量的名字 | 命名要规范，后面TensorRT要用 |
output_names |
输出张量的名字 | 同上，建议用有意义的名称 |
dynamic_axes |
哪些维度是动态的 | YOLO推理时batch_size和宽高经常变 |
torch.jit.trace或torch.jit.script试试，能提前发现一些图结构问题。我遇到过好几次模型里有动态控制流，直接导出ONNX会报错。
3.2 动态轴与静态轴设置
这是最容易出问题的地方。静态轴，就是导出时固定死的维度。动态轴，就是推理时可以变化的维度。
举个例子。YOLO模型输入是[batch, 3, height, width]。如果你导出时固定了batch=1，那推理时只能一次处理一张图。但实际场景里，你可能想一次处理4张、8张甚至16张。这时候就需要把batch维度设成动态的。
动态轴的配置方式是这样的：
dynamic_axes = {
'input': {
0: 'batch_size', # 第0维是batch，名字叫batch_size
2: 'height', # 第2维是高度
3: 'width' # 第3维是宽度
},
'output': {
0: 'batch_size' # 输出也要对应动态
}
}
这里有个坑。YOLO的输出通常包含多个head，比如检测头、分割头。每个输出的形状可能不一样。我建议只把batch维度设成动态，宽高尽量固定。为什么呢？
- 宽高动态：会导致输出特征图的尺寸变化，后续的NMS处理逻辑要跟着变，容易出bug。
- batch动态：对输出形状影响小，只是多几张图的数据堆在一起。
那什么时候必须用动态宽高？比如你的应用需要处理不同分辨率的输入，像手机拍照的尺寸不固定。这时候可以设，但要做好性能下降的心理准备。
3.3 输入输出命名规范
命名这事，看着不起眼，但真能坑人。我刚开始做部署时，随便起了个名字叫input1、output1，结果后面TensorRT绑定数据时，半天找不到对应的张量。
我的命名习惯是这样的：
- 输入：用
input或images，简单明了。 - 输出：用
output或detections，如果是多输出，用output_0、output_1。 - 中间张量：一般不需要命名，除非你要做中间层调试。
命名时注意几点：
- 不要用特殊字符：比如空格、斜杠、点号。有些推理引擎解析时会报错。
- 不要用中文：ONNX标准里不支持中文张量名。
- 保持一致性：输入输出名字要和后续代码里用的一致。
经验之谈： 我建议在项目里统一一个命名规范文档。比如输入都叫input，输出都叫output，中间层用layer_xxx。这样团队协作时，每个人都知道该绑哪个张量。
3.4 opset版本选择
opset，就是ONNX算子集的版本号。每个版本会新增或废弃一些算子。选错了，模型可能导出失败，或者推理时某些算子不支持。
YOLO模型常用的opset版本：
| opset版本 | 特点 | YOLO兼容性 |
|---|---|---|
| 10 | 基础算子，支持较少 | YOLOv5早期版本可用，但容易报错 |
| 11 | 增加了ScatterND、GatherElements等 |
YOLOv5/v8推荐，兼容性好 |
| 12 | 增加了ReduceL1、ReduceL2等 |
YOLOv8可用，但部分算子可能不被TensorRT支持 |
| 13 | 增加了SoftmaxCrossEntropyLoss等 |
YOLOv8可用，但建议先测试TensorRT兼容性 |
| 14+ | 新算子多，但兼容性风险大 | 不推荐，除非你确定推理引擎支持 |
我个人建议：YOLO模型导出ONNX时，opset选11或12。这两个版本在TensorRT上支持得最好。选太高了，TensorRT可能还没实现那些新算子，推理时会fallback到CPU，性能直接崩掉。
ScatterND算子。结果TensorRT 8.4不支持，只能跑在CPU上，推理速度从30ms掉到300ms。后来改成opset=11，一切正常。
3.5 知识结构图：ONNX导出核心流程
下面这张图，把整个导出流程串起来了。你可以对照着看，每一步该做什么。
3.6 验证导出结果
导出完ONNX，别急着往下走。先验证一下模型对不对。我一般用ONNX Runtime跑一遍，看看输出和PyTorch原模型是否一致。
import onnx
import onnxruntime as ort
import numpy as np
# 加载ONNX模型
onnx_model = onnx.load('yolov8n.onnx')
onnx.checker.check_model(onnx_model) # 检查模型结构
# 创建推理会话
ort_session = ort.InferenceSession('yolov8n.onnx')
# 准备输入数据
input_name = ort_session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 640, 640).astype(np.float32)
# 运行推理
outputs = ort_session.run(None, {input_name: input_data})
# 对比PyTorch输出
with torch.no_grad():
torch_outputs = model(torch.from_numpy(input_data))
# 检查数值差异
for i, (ort_out, torch_out) in enumerate(zip(outputs, torch_outputs)):
diff = np.abs(ort_out - torch_out.numpy()).max()
print(f'输出 {i} 最大差异: {diff:.6f}')
if diff < 1e-3:
print('✅ 验证通过')
else:
print('❌ 差异过大，请检查模型')
这里有个细节。ONNX Runtime和PyTorch的数值精度不完全一样，因为算子实现有差异。我一般容忍1e-3的误差。如果差异超过这个值，就要排查是不是导出时某些算子被替换了。
onnxruntime.transformers里的工具，把ONNX模型可视化，看看计算图有没有异常。我遇到过几次，模型里多了些奇怪的Identity节点，删掉就好了。
嗯，这一章的内容就这些。torch.onnx.export看着简单，但每个参数都值得仔细推敲。动态轴设好了，后面TensorRT优化事半功倍。命名规范了，团队协作不吵架。opset选对了，少踩一半的坑。