4、ONNX基础：ONNX格式介绍、如何将PyTorch模型导出为ONNX、ONNX opset版本选择
好，咱们进入第四章。这一章聊ONNX，它是TensorRT生态里绕不开的一环。说白了，ONNX就是模型界的“通用语言”。你拿PyTorch训练好的模型，想部署到TensorRT上跑，中间总得有个桥梁。ONNX就是这个桥梁。
我个人习惯，不管最终用TensorRT还是ONNX Runtime，都会先把模型过一遍ONNX。为什么？因为ONNX能帮你提前发现很多算子兼容性问题。你想想看，在PyTorch里跑得飞起的模型，导出成ONNX可能就报错了。早发现早解决，总比到了部署阶段再抓瞎强。
4.1 ONNX格式到底是个啥？
ONNX全称Open Neural Network Exchange，微软和Facebook联合搞的。它的核心思想很简单：定义一个中间表示（IR），让不同框架的模型都能转成这个格式。
ONNX模型文件本质上是一个protobuf序列化的二进制文件。你用onnx.load()读进来，就是一个ModelProto对象。这个对象里主要包含三部分：
- GraphProto：计算图本身，包含所有节点（NodeProto）、输入输出（ValueInfoProto）、初始值（TensorProto）
- MetadataProps：模型的元信息，比如作者、版本、opset版本等
- ModelProto：最外层容器，包含ir_version、producer_name等
嗯，这里要注意：ONNX的图是静态图。什么意思？就是你的模型结构在导出那一刻就固定了，不能再动态改变。比如PyTorch里那种if x.shape[0] > 10: ... else: ...的动态分支，ONNX是没法直接表达的。我在项目中遇到过好几次这种问题，后面会讲怎么处理。
核心要点：ONNX不是运行时，它只是一个模型描述格式。它定义了算子（Operator）和数据类型，但不负责怎么执行。执行是TensorRT、ONNX Runtime这些推理引擎的事。
4.2 如何把PyTorch模型导出为ONNX？
PyTorch从1.2版本开始就内置了torch.onnx.export()函数。用法其实不复杂，但坑不少。我给你看个标准流程：
import torch
import torch.onnx
# 假设你有一个训练好的模型
model = MyResNet50()
model.load_state_dict(torch.load('model.pth'))
model.eval()
# 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 输入张量
'resnet50.onnx', # 输出文件名
export_params=True, # 是否导出参数权重
opset_version=11, # opset版本，后面细说
do_constant_folding=True, # 常量折叠优化
input_names=['input'], # 输入节点名称
output_names=['output'], # 输出节点名称
dynamic_axes={
'input': {0: 'batch_size'}, # 动态batch
'output': {0: 'batch_size'}
}
)这段代码我用了无数次。有几个地方我想特别提醒你：
- model.eval()一定要加。我见过有人忘了这步，结果导出的ONNX里包含了Dropout和BatchNorm的训练行为，推理结果完全不对。
- dummy_input的shape要和实际输入一致。特别是batch维度，如果你后面要用动态batch，这里可以设成1，但shape的其他维度必须准确。
- dynamic_axes参数：这个很关键。它告诉ONNX哪些维度是动态的。比如
{0: 'batch_size'}表示batch维度可变。TensorRT处理动态batch时，就是靠这个信息来构建优化配置文件的。
个人经验：导出ONNX时，我建议先用torch.onnx.export()的verbose=True参数跑一遍。它会打印出每个节点的详细信息，方便你排查算子映射问题。另外，导出后用onnx.checker.check_model()验证一下，能发现很多低级错误。
4.3 ONNX opset版本选择——这个坑我踩过
opset版本，说白了就是ONNX算子规范的版本号。每个版本会新增、修改或废弃一些算子。比如opset 11引入了ScatterND，opset 13改进了Softmax的行为。
那到底选哪个版本？我的建议是：
| 场景 | 推荐opset版本 | 原因 |
|---|---|---|
| 通用部署（TensorRT 8.x） | 11 或 13 | 兼容性好，TensorRT支持完善 |
| 需要动态shape | 11+ | opset 11开始对动态shape支持更好 |
| 使用Transformer/BERT | 13 或 15 | 新版本对Attention类算子有优化 |
| TensorRT 10.x最新版 | 17 或 18 | 支持更多新算子，但注意兼容性 |
我曾经在项目里选了个过高的opset版本（比如opset 18），结果TensorRT解析时直接报“Unsupported operator”。后来查文档才发现，TensorRT对opset的支持是滞后的。你想想看，ONNX每年发两个版本，TensorRT不可能同步跟上。
避坑指南：我曾经因为opset版本选错，导致模型导出后某些算子被拆成了多个小算子，推理性能下降了30%。后来我固定用opset 11，问题就解决了。所以我的建议是：除非你明确需要某个新版本的功能，否则就用opset 11。它稳定、兼容性好，TensorRT支持得最全面。
还有一个细节：PyTorch的torch.onnx.export()默认opset版本是9（老版本PyTorch）或11（新版本）。你可以通过torch.onnx.supported_ops()查看当前PyTorch支持哪些opset版本。
4.4 导出后的验证与调试
模型导出完，别急着扔给TensorRT。先做两件事：
- 数值一致性检查：用ONNX Runtime跑一遍推理，对比PyTorch的输出。误差应该在1e-5以内。如果偏差大，说明导出过程中有精度损失或算子映射问题。
- 可视化检查：用
netron（一个在线工具）打开ONNX文件，看看图结构对不对。我习惯检查输入输出节点的名称和shape，还有关键算子的连接是否正确。
import onnx
import onnxruntime as ort
import numpy as np
# 加载ONNX模型
onnx_model = onnx.load('resnet50.onnx')
onnx.checker.check_model(onnx_model)
# 用ONNX Runtime推理
ort_session = ort.InferenceSession('resnet50.onnx')
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)
# 对比PyTorch输出
with torch.no_grad():
torch_output = model(dummy_input)
print('Max difference:', np.max(np.abs(ort_outputs[0] - torch_output.numpy())))嗯，这一步虽然简单，但能省掉你后面很多调试时间。我每次导出ONNX都会跑一遍这个检查，已经成习惯了。
4.5 小结
ONNX是PyTorch到TensorRT的必经之路。记住三个要点：
- 导出时用
model.eval()，设对dynamic_axes - opset版本选11最稳妥，除非你有特殊需求
- 导出后一定要做数值一致性检查
下一章我们会讲TensorRT如何解析ONNX模型，以及那些常见的解析失败问题怎么解决。到时候你会看到，很多坑其实在ONNX导出阶段就能避免。