3、模型序列化基础：ONNX 导出要点
好，咱们进入第三章。模型序列化，说白了就是把训练好的模型“打包”成一个标准格式。在 TensorRT 的世界里，ONNX 就是这个标准格式的“通用语言”。
我个人习惯把这一步叫做“模型翻译”。你想想看，PyTorch 模型就像一本中文书，TensorRT 只读英文。ONNX 就是那个翻译官。翻译得好不好，直接决定了后面 TensorRT 优化的天花板。
3.1 ONNX 导出要点：别踩这些坑
先说说最基本的导出流程。PyTorch 转 ONNX，核心就一个函数：torch.onnx.export()。但参数怎么设，这里头门道不少。
import torch
import torchvision.models as models
# 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()
# 构造一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)
# 导出 ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 示例输入
"resnet18.onnx", # 输出文件名
export_params=True, # 导出参数
opset_version=17, # opset 版本，我建议用 17 或更高
do_constant_folding=True, # 常量折叠优化
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
嗯，这里有几个要点我得重点说说：
- opset_version：这个版本号决定了 ONNX 支持哪些算子。版本太低，很多新算子不支持；版本太高，可能某些推理引擎不兼容。我个人建议用 17 或 18，TensorRT 8.x 和 10.x 都支持得很好。
- dynamic_axes：如果你需要动态 batch 或者动态输入尺寸，一定要设置这个。不设置的话，导出的 ONNX 就固定死了 batch size=1，后面 TensorRT 优化时灵活性大打折扣。
- do_constant_folding：这个建议打开。它会把一些固定的计算提前算好，比如 BN 层的融合。能减小模型体积，也能加速推理。
model.eval()。我见过不少新手忘了这步，结果导出的 ONNX 里还带着 dropout 和 BN 的训练模式，推理结果完全不对。
3.2 PyTorch 模型转 ONNX 的常见陷阱
说实话，PyTorch 转 ONNX 的过程，十次有八次不会一次成功。我在项目中遇到过各种奇奇怪怪的问题，下面这几个是最常见的。
陷阱一：动态控制流
PyTorch 模型里如果有 if 语句或者 for 循环，而且这些控制流依赖于输入数据，那 ONNX 导出就会报错。为什么？因为 ONNX 是静态图，它要求计算路径是固定的。
# 错误示例：动态控制流
class BadModel(torch.nn.Module):
def forward(self, x):
if x.sum() > 0: # 这个 if 依赖于输入数据
return x * 2
else:
return x * 3
# 正确做法：用 ONNX 支持的算子替代
class GoodModel(torch.nn.Module):
def forward(self, x):
# 用 torch.where 替代 if-else
condition = (x.sum() > 0).float()
return x * (2 + condition)
我曾经接手过一个项目，模型里有个循环做序列处理，导出 ONNX 时直接报“不支持动态循环”。最后只能把循环展开成固定长度的计算，或者用 ONNX 的 Loop 算子重写。
陷阱二：不支持的操作
PyTorch 有上千个算子，但 ONNX 支持的只有几百个。像 torch.einsum、torch.topk 的某些模式、自定义的 CUDA 扩展，这些在导出时都可能翻车。
| PyTorch 操作 | ONNX 支持情况 | 替代方案 |
|---|---|---|
| torch.einsum | 部分支持（opset >= 16） | 用 matmul + reshape 组合 |
| torch.topk | 支持（opset >= 11） | 直接使用即可 |
| F.interpolate | 支持，但模式有限 | 用 'nearest' 或 'bilinear' |
| 自定义 CUDA 算子 | 不支持 | 用 ONNX 原生算子重写 |
陷阱三：输入输出形状不匹配
这个坑我踩过好几次。导出的 ONNX 在 PyTorch 里验证没问题，但一拿到 TensorRT 里就报形状错误。原因往往是模型里有 reshape 或者 view 操作，它们的输出形状依赖于输入数据。
# 问题代码：动态 reshape
class ReshapeModel(torch.nn.Module):
def forward(self, x):
batch_size = x.shape[0]
# 这里 reshape 的形状依赖于 batch_size
return x.view(batch_size, -1, 4)
# 解决方案：用 ONNX 的 Reshape 算子
# 或者固定形状，避免动态依赖
嗯，这里有个调试技巧：导出 ONNX 后，用 onnxruntime 先跑一遍推理，看看结果对不对。如果 onnxruntime 能跑通，那 TensorRT 那边大概率也没问题。
3.3 调试方法：三板斧
遇到 ONNX 导出问题，别慌。我总结了一套调试方法，管用。
- 第一板斧：简化模型。把模型拆成小块，逐块导出。比如先导 backbone，再导 head。哪块出问题，一目了然。
- 第二板斧：用 onnx.checker 验证。导出后立刻跑一下检查，能发现结构性问题。
- 第三板斧：可视化。用 Netron 打开 ONNX 文件，看看计算图是不是你想要的。我经常发现有些算子被莫名其妙地拆成了好几个小算子，这就是优化的线索。
import onnx
# 加载 ONNX 模型
onnx_model = onnx.load("resnet18.onnx")
# 检查模型结构
onnx.checker.check_model(onnx_model)
# 打印模型信息
print(onnx.helper.printable_graph(onnx_model.graph))
最后说一句，如果你在导出时遇到“Unsupported operator”的错误，别急着改模型。先去 ONNX 官方算子列表查一下，看看是不是 opset 版本太低。我遇到过好几次，升级 opset 版本就解决了。
好，这一章就到这里。下一章我们聊聊 ONNX 的图优化技巧，怎么把导出的 ONNX 模型“瘦身”和“整形”，让它更适合 TensorRT 的胃口。