📚 课程目录第3章模型序列化基础

3. 模型序列化基础：ONNX 中间格式与导出优化

模型训练好之后，下一步就是部署。但 PyTorch 模型不能直接跑在 TensorRT 上，中间需要一座桥——ONNX。说白了，ONNX 就是个模型交换格式，让不同框架能互相理解。

我个人习惯把 ONNX 看作模型的「通用语言」。你想想看，PyTorch 训练出来的模型，要跑到 NVIDIA 的 TensorRT 上，中间没有 ONNX 做翻译，根本行不通。这一章我们就来搞定这个翻译过程。

  核心要点：ONNX 是模型序列化的中间桥梁，torch.onnx.export 是 PyTorch 到 ONNX 的标准工具，ONNX-Simplifier 负责清理冗余结构。

3.1 为什么需要 ONNX？

ONNX（Open Neural Network Exchange）由微软和 Facebook 联合推出。它的目标很明确：让模型在不同框架之间自由迁移。

我遇到过不少团队，训练用 PyTorch，推理想用 TensorRT，中间卡在模型转换上。没有 ONNX 的话，你得手写网络结构再加载权重，那工作量……嗯，想想就头疼。

特性	PyTorch 模型	ONNX 模型	TensorRT 引擎
格式	.pt / .pth	.onnx	.engine / .trt
依赖框架	PyTorch 运行时	ONNX Runtime / 任意框架	TensorRT 运行时
推理速度	一般	中等	极快
部署灵活性	低	高	最高

ONNX 的核心价值在于：一次导出，多处部署。你只需要把 PyTorch 模型转成 ONNX，后续无论是 TensorRT、ONNX Runtime 还是其他推理引擎，都能直接使用。

3.2 torch.onnx.export 实战

PyTorch 官方提供了 torch.onnx.export 函数，专门用来把模型导出为 ONNX 格式。我们来看一个完整的例子。

3.2.1 基本用法

import torch
import torchvision.models as models

# 1. 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 2. 创建 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出 ONNX
torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 示例输入
    "resnet18.onnx",          # 输出文件名
    export_params=True,       # 导出参数
    opset_version=11,         # ONNX opset 版本
    do_constant_folding=True, # 常量折叠优化
    input_names=['input'],    # 输入名称
    output_names=['output'],  # 输出名称
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

print("ONNX 模型导出成功！")

这里有几个关键参数，我重点说一下：

opset_version：ONNX 算子集版本。我建议用 11 或更高版本，支持更多算子。
do_constant_folding：把常量计算提前算好，减少推理时的计算量。
dynamic_axes：允许动态 batch size。我在项目中经常遇到推理时 batch 大小不固定的情况，这个参数非常有用。

小技巧：导出前一定要调用 model.eval()，关闭 dropout 和 batch normalization 的训练行为。我曾经忘了这步，导出的 ONNX 模型推理结果一直不对，排查了半天才发现是 dropout 还在生效。

3.2.2 验证 ONNX 模型

导出之后，最好验证一下 ONNX 模型是否正确。用 ONNX Runtime 跑一遍推理，对比 PyTorch 的输出。

import onnx
import onnxruntime as ort
import numpy as np

# 检查 ONNX 模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型结构检查通过！")

# 用 ONNX Runtime 推理
ort_session = ort.InferenceSession("resnet18.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)

# 检查误差
diff = np.abs(ort_outputs[0] - torch_output.numpy()).max()
print(f"最大误差：{diff:.6f}")

if diff < 1e-4:
    print("✅ 模型导出正确！")
else:
    print("⚠️ 误差较大，请检查模型")

为什么误差要控制在 1e-4 以内？因为浮点数计算本身就有精度损失，这个范围内的误差是可以接受的。如果误差太大，说明导出过程有问题。

3.3 ONNX-Simplifier 优化模型结构

导出的 ONNX 模型往往包含很多冗余结构。比如 PyTorch 自动插入的 Shape 节点、Cast 节点，还有常量折叠后留下的死节点。这些冗余不会影响精度，但会增加模型大小和推理延迟。

ONNX-Simplifier 就是用来干这个的。它会分析计算图，删除无用节点，合并可以简化的结构。

3.3.1 安装与使用

# 安装
pip install onnx-simplifier

# 命令行使用
python -m onnxsim resnet18.onnx resnet18_simplified.onnx

# Python API 使用
import onnx
from onnxsim import simplify

model = onnx.load("resnet18.onnx")
model_simplified, check = simplify(model)

if check:
    onnx.save(model_simplified, "resnet18_simplified.onnx")
    print("简化成功！")
else:
    print("简化失败，请检查模型")

3.3.2 简化效果对比

指标	简化前	简化后	优化比例
节点数	176	142	19.3%
模型大小	44.7 MB	44.7 MB	0%
推理延迟	12.3 ms	11.8 ms	4.1%

注意看，模型大小基本没变，因为权重参数没少。但节点数减少了 19%，推理延迟也降了 4%。别小看这 4%，在边缘设备上，每一毫秒都很宝贵。

避坑指南：我曾经遇到过一个模型，简化后推理结果完全不对。排查发现是 ONNX-Simplifier 把某些自定义算子当成了冗余节点删掉了。所以简化后一定要做精度验证，不要盲目信任工具。

3.4 知识体系总览

下面这张图展示了模型序列化的完整流程，从 PyTorch 到 ONNX，再到简化优化，最后输入 TensorRT。

3.5 常见问题与避坑

做模型序列化这几年，我踩过不少坑。挑几个典型的说说：

动态 shape 问题：有些算子不支持动态 shape，导出时会报错。我建议先用固定 shape 导出，确认没问题后再加 dynamic_axes。
自定义算子：如果你的模型用了自定义算子，ONNX 可能不认识。需要注册自定义算子，或者用替代方案。
版本兼容性：PyTorch、ONNX、ONNX-Simplifier 的版本要匹配。我遇到过 PyTorch 1.10 导出的模型，ONNX Runtime 1.8 跑不了的情况。

我的经验：每次导出 ONNX 后，我都会用 onnx.shape_inference.infer_shapes 做一次 shape 推理，确保所有节点的输入输出 shape 都正确。这一步能提前发现很多问题。

好了，模型序列化这块就讲到这里。你只要掌握了 torch.onnx.export 和 ONNX-Simplifier，后续的 TensorRT 转换就会顺畅很多。记住：ONNX 质量决定了 TensorRT 的上限，花时间把 ONNX 模型优化好，后面事半功倍。