3. ONNX 导出原理：从 PyTorch 到通用中间格式

好，咱们进入第三章。这一章我打算把 ONNX 导出这件事彻底讲透。很多同学跑通了导出，但遇到动态尺寸报错、算子不兼容、精度对不上就懵了。说白了，就是没搞懂 ONNX 到底是个什么东西，以及 torch.onnx.export 背后在干什么。

我个人习惯把 ONNX 看作「模型的世界语」。你训练好的 PyTorch 模型，别人可能用 TensorFlow、用 Caffe、用 CoreML，怎么互通？ONNX 就是那个中间桥梁。而 TensorRT 部署，只是这座桥梁的一个目的地。

3.1 ONNX 格式：不只是个「模型文件」

ONNX（Open Neural Network Exchange）本质上是一个计算图的表示协议。它不关心你怎么训练的，只关心网络长什么样、数据怎么流动。

一个 ONNX 模型文件（.onnx）内部包含三样东西：

  • 计算图结构：节点（算子）和边（张量）构成的有向无环图
  • 权重参数：所有可训练参数序列化存储
  • 元数据：输入输出名称、数据类型、形状信息、opset 版本等

你可以用 onnx.load() 把它读进来，然后用 onnx.checker.check_model() 验证完整性。我在项目中经常遇到导出时没报错，但 checker 一跑就发现图结构有问题的案例——嗯，这一步千万别省。

核心认知：ONNX 是静态图。PyTorch 是动态图（define-by-run），导出时 torch.onnx.export 会 trace 一遍 forward，把动态执行路径「凝固」成静态计算图。这意味着——控制流（if/for）默认不会被导出，除非你用 torch.jit.script 或显式标注。

3.2 torch.onnx.export 参数详解：每个参数我都踩过坑

这个函数是导出流程的核心。我把它最常用的参数列出来，每个都附上我在实战中遇到的教训。

参数 类型 说明 避坑指南
model torch.nn.Module 要导出的模型实例 记得调成 model.eval() 模式，否则 BN/Dropout 行为会错
args tuple 模型的示例输入 形状和数据类型必须与真实推理一致，否则 trace 出来的图是错的
f str 输出文件路径 建议用 .onnx 后缀，不要用 .pt 或 .pth
export_params bool 是否导出权重 默认 True。部署场景必须 True，否则推理结果全错
opset_version int ONNX 算子集版本 YOLO 系列建议 11~17，后面单独讲
do_constant_folding bool 是否做常量折叠优化 建议 True，能减少冗余计算节点
input_names list[str] 输入张量名称 YOLO 通常用 ['images']
output_names list[str] 输出张量名称 YOLOv8 用 ['output0']，注意区分不同版本
dynamic_axes dict 定义动态轴 YOLO 部署必配，后面重点讲

我曾经踩过的坑：有一次导出 YOLOv5，忘记设 model.eval()，结果导出的 ONNX 在 TensorRT 上推理结果全是 NaN。查了两天才发现是 BN 层在 training 模式下统计了错误的 running mean。所以——导出前先调 eval，这是铁律。

3.3 动态轴与静态轴：YOLO 部署的命门

这是整个导出环节最容易出错的地方。我先说清楚概念。

静态轴：导出时固定死的维度。比如 batch=1，那导出的 ONNX 就只能一次推理一张图。想改 batch size？重新导出。

动态轴：导出时不固定，推理时再指定。比如 batch 可以是 1、4、8、16，同一个 ONNX 文件都能处理。

YOLO 部署场景下，我们通常需要动态的 batch 和动态的 height/width。为什么？因为生产环境里，你不可能保证每次请求的图片尺寸都一样。

来看一个 YOLOv8 的典型配置：

dynamic_axes = {
    'images': {
        0: 'batch',
        2: 'height',
        3: 'width'
    },
    'output0': {
        0: 'batch',
        2: 'height',
        3: 'width'
    }
}

torch.onnx.export(
    model,
    dummy_input,
    'yolov8n.onnx',
    input_names=['images'],
    output_names=['output0'],
    dynamic_axes=dynamic_axes,
    opset_version=17,
    do_constant_folding=True
)

注意看：images 的形状是 (batch, 3, height, width)，我把第 0、2、3 维都设成了动态。第 1 维是通道数（RGB=3），这个通常是固定的，不需要动。

我的建议：如果你只做推理服务，batch 维度设动态就够了。height 和 width 设动态会引入额外的 resize 操作，在 TensorRT 上可能带来性能损失。我一般做法是：固定输入尺寸（如 640x640），只动态 batch。这样既灵活又高效。

3.4 opset 版本选择：不是越新越好

opset 就是 ONNX 算子集的版本号。每个版本会新增、废弃或修改一些算子。选错了，要么导出报错，要么推理结果不对。

我整理了一个版本选择参考表：

opset 版本 关键变化 YOLO 适配建议
11 支持 Resize 的 coordinate_transformation_mode YOLOv5 最低要求
13 支持 Softmax 的 axis 参数 YOLOv8 常用基线
15 支持 Shape 的 start/end 切片 部分后处理算子需要
17 支持 LayerNormalization 等新算子 YOLOv8/v11 推荐
18+ 新增 GroupNormalization 等 TensorRT 可能不支持，慎用

我个人习惯：YOLOv5 用 opset 12，YOLOv8/v11 用 opset 17。为什么不用 18？因为 TensorRT 8.x 对 opset 18 的支持还不完善，我在项目里遇到过 GroupNormalization 算子无法转换的问题。你想想看，导出成功了，但 TensorRT 编译失败，那不等于白干？

注意：opset 版本不是越高越好。高版本算子可能不被目标推理引擎支持。选版本前，先查一下你的 TensorRT 版本支持的最高 opset。比如 TensorRT 8.5 最高支持 opset 17，你选 18 就会报 Unsupported operator。

3.5 完整导出流程：从 PyTorch 到 ONNX

把上面所有知识点串起来，一个完整的 YOLO 导出流程长这样：

  1. 加载模型：torch.load() 或 ultralytics.YOLO()
  2. 调 eval 模式：model.eval()
  3. 构造 dummy input：形状与真实输入一致，值可以是随机数
  4. 配置 dynamic_axes：至少动态 batch
  5. 选择 opset：YOLO 推荐 12~17
  6. 执行导出：torch.onnx.export()
  7. 验证 ONNX：onnx.checker.check_model()
  8. 测试推理：用 onnxruntime 跑一遍，对比 PyTorch 输出

最后一步很多人会跳过。但我建议你务必做——用 ONNX Runtime 跑一次推理，把输出和 PyTorch 的原始输出做数值对比。误差在 1e-3 以内算正常，超过 1e-2 就要排查了。

经验之谈：我遇到过导出时一切正常，但 ONNX Runtime 推理结果和 PyTorch 差了 0.5 的情况。最后发现是 torch.onnx.export 的 do_constant_folding 把某些 BN 层的参数折叠错了。关掉这个选项就恢复了。所以——不要盲目信任默认参数，每个选项都要理解它的影响。

3.6 知识结构图：ONNX 导出全景

下面这张图把本章的核心知识点串起来了。你可以把它当作导出前的检查清单。

ONNX 导出全景图 PyTorch 模型 torch.onnx.export() 参数配置 input_names / output_names 动态轴 batch / height / width opset 版本 11 / 13 / 15 / 17 model.onnx 验证：checker + ONNX Runtime 推理

这张图从左到右展示了导出流程的三个关键决策点：参数怎么配、动态轴怎么设、opset 怎么选。每一步都直接影响最终 ONNX 文件的质量。我建议你把它打印出来贴在工位上，每次导出前扫一眼，能少踩很多坑。

← 上一章 📖 返回目录 下一章 →