3、ONNX基础：ONNX格式详解、ONNX算子标准、PyTorch模型导出为ONNX

这一章，我们来啃ONNX这块硬骨头。

说实话，我刚开始接触ONNX时也觉得它挺玄乎的。一个中间格式，凭什么能打通那么多框架？后来踩了不少坑，才慢慢摸清楚它的脾气。今天我把这些经验整理出来，希望能帮你少走弯路。

3.1 ONNX格式到底长什么样？

ONNX的全称是Open Neural Network Exchange。说白了，它就是一套标准化的模型描述语言。不管你是用PyTorch、TensorFlow还是其他框架训练出来的模型，都能转换成ONNX格式。

那ONNX文件里到底装了啥？我直接给你拆开看。

一个ONNX模型的核心结构，其实就三大部分：

• 计算图（Graph）：描述网络的结构，包括节点、边、张量
• 算子集（Opset）：定义每个节点用什么操作，比如Conv、Relu、MatMul
• 元数据（Metadata）：模型名称、版本、作者、文档等信息

你可以把ONNX想象成一张蓝图。蓝图本身不干活，但它精确描述了每一块砖应该怎么放。TensorRT拿到这张蓝图后，就知道该怎么优化和执行。

下面这张图，是我画的一个ONNX模型内部结构示意，帮你建立直观印象：

3.2 ONNX算子标准——你不得不知道的事

ONNX定义了一套标准算子，比如Conv、Relu、BatchNormalization、Gemm（全连接）等等。每个算子都有固定的输入输出格式和属性参数。

这里有个关键点：不同框架的算子，映射到ONNX时可能不完全一致。比如PyTorch的padding模式，和ONNX的padding模式就有细微差别。我在项目中就遇到过，一个Conv2d导出后，因为padding参数没对齐，推理结果差了0.01。

下面这个表格，是我整理的常用ONNX算子与PyTorch算子的对应关系：

PyTorch 算子	ONNX 算子	注意事项
torch.nn.Conv2d	Conv	padding模式需对齐，dilation>1时注意
torch.nn.ReLU	Relu	基本一致，无坑
torch.nn.BatchNorm2d	BatchNormalization	training=False时导出，否则多出训练节点
torch.nn.Linear	Gemm	转置权重时注意transA/transB参数
torch.nn.MaxPool2d	MaxPool	ceil_mode需手动指定
torch.nn.AdaptiveAvgPool2d	GlobalAveragePool	输出尺寸为1x1时自动映射
torch.cat	Concat	axis参数从0开始，与PyTorch一致
torch.reshape	Reshape	允许-1维度，但ONNX要求shape是常量

3.3 PyTorch模型导出为ONNX——torch.onnx.export实战

好，理论说完了，咱们直接上手。PyTorch导出ONNX的核心函数就一个：torch.onnx.export()。但参数不少，我一个个给你拆。

先看一个最基础的导出示例：

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = torch.load('my_model.pth')
model.eval()  # 一定要切到eval模式

# 构造一个 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出
torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 示例输入
    'my_model.onnx',          # 输出路径
    export_params=True,       # 导出训练好的参数
    opset_version=11,         # 算子集版本
    do_constant_folding=True, # 常量折叠优化
    input_names=['input'],    # 输入节点名称
    output_names=['output'],  # 输出节点名称
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)
print("导出完成！")

这段代码看着简单，但每个参数都有讲究。我挑几个重点说说：

• opset_version：我建议用11或更高。版本太低，有些新算子不支持；版本太高，老版本TensorRT可能不兼容。我个人习惯用opset=11，兼容性最好。
• dynamic_axes：这个参数很关键。它告诉ONNX哪些维度是可变的。比如batch_size，你训练时是1，推理时可能是8或16。不设置的话，导出的模型就固定batch_size=1了。
• do_constant_folding：建议开启。它会提前计算一些常量表达式，减小模型体积，提升推理速度。

3.4 导出后的验证——别急着往下走

导出完ONNX文件，第一件事不是扔给TensorRT，而是先验证一下。

我一般做三步验证：

1. 结构验证：用 onnx.checker.check_model() 检查模型结构是否合法
2. 精度验证：用onnxruntime跑一遍，对比PyTorch的输出，误差应该在1e-5以内
3. 可视化验证：用Netron打开ONNX文件，看看网络结构对不对

下面这段代码，是我常用的验证流程：

import onnx
import onnxruntime as ort
import numpy as np

# 1. 结构验证
onnx_model = onnx.load('my_model.onnx')
onnx.checker.check_model(onnx_model)
print("✅ 模型结构验证通过")

# 2. 精度验证
ort_session = ort.InferenceSession('my_model.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).numpy()

diff = np.abs(ort_outputs[0] - torch_output).max()
print(f"✅ 最大误差: {diff:.2e}")
if diff < 1e-5:
    print("✅ 精度验证通过")
else:
    print("⚠️ 精度偏差较大，请检查")

3.5 常见导出问题与解决方案

导出ONNX时，你大概率会遇到下面这些问题。我直接给你列出来，省得你再去搜：

问题现象	根本原因	解决方案
导出时报错"Unsupported operator"	PyTorch算子不在ONNX算子集中	用torch.onnx.register_custom_op_symbolic注册，或换用ONNX支持的实现
导出成功但推理结果不对	模型未切到eval模式，或算子参数不匹配	检查model.eval()，对比每个算子的参数
动态batch导出失败	dynamic_axes设置错误，或模型中有固定shape的操作	检查所有reshape/view操作，确保batch维度是动态的
ONNX文件体积过大	包含了不必要的常量或调试信息	开启do_constant_folding，导出后可以用onnx-simplifier简化
TensorRT加载ONNX报错	算子版本不兼容，或使用了TensorRT不支持的算子	降低opset版本，或替换算子为TensorRT支持的版本