第2章 ONNX中间表示

第2章：ONNX中间表示：ONNX的由来与优势，PyTorch模型转ONNX的标准流程

2.1 ONNX是什么？为什么我们需要它？

说实话，我刚入行做模型部署那会儿，最头疼的就是框架之间的转换。你想想看，PyTorch训练好的模型，想放到TensorRT上跑，或者想用ONNX Runtime做推理，中间得折腾多少事？

ONNX（Open Neural Network Exchange）就是来解决这个问题的。它像一个“通用语言”，让不同深度学习框架能互相沟通。说白了，ONNX定义了一套标准的计算图表示方法。不管你是用PyTorch、TensorFlow还是其他框架，只要导出成ONNX格式，就能在各种推理引擎上跑起来。

我在项目中遇到过最典型的场景：客户要求模型必须用TensorRT做推理加速，但团队里所有人都只会PyTorch。嗯，这时候ONNX就成了我们的“救星”。

ONNX的核心优势：

框架无关性：一次导出，到处运行
算子标准化：定义了统一的算子集，避免各家实现差异
生态丰富：几乎所有主流推理引擎都支持ONNX
图优化能力：ONNX本身支持常量折叠、算子融合等优化

2.2 PyTorch转ONNX的标准流程

好，我们直接进入正题。PyTorch转ONNX，核心就一个函数：torch.onnx.export()。但别小看它，里面坑不少。

我个人习惯把导出流程分成三步：

准备模型和输入：确保模型处于eval模式，准备一个dummy输入
配置导出参数：设置输入输出名称、动态轴、opset版本等
验证导出的ONNX模型：用onnxruntime跑一下，对比精度

来看一个最基础的例子：

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",
    input_names=["input"],
    output_names=["output"],
    opset_version=11,
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

print("导出成功！")

小提示：opset_version建议用11或更高版本。我刚开始用opset=9，结果有些算子不支持，折腾了半天。现在一般用11或13，兼容性最好。

2.3 关键参数详解

上面代码里那几个参数，每个都有讲究。我一个个说：

参数	作用	我的建议
`input_names`	定义输入节点的名称	用有意义的名称，比如"input"、"image"
`output_names`	定义输出节点的名称	同理，用"output"、"logits"等
`opset_version`	指定ONNX算子集版本	推荐11或13，别用太老的
`dynamic_axes`	指定哪些维度是动态的	batch_size几乎必设，其他看需求

你可能会问：为什么要有dynamic_axes？

举个例子。你训练时batch_size是32，但部署时可能一次只推理1张图。如果不设置动态轴，导出的ONNX就固定死了batch_size=32。到时候你传个batch=1的输入，直接报错。

注意：动态轴虽然灵活，但有些推理引擎对动态shape支持不好。比如TensorRT，动态batch需要额外配置。我的经验是：如果业务场景固定，尽量用静态shape，性能更好。

2.4 常见坑与避坑指南

导出ONNX看着简单，但实际项目中我踩过不少坑。分享几个典型的：

算子不支持：PyTorch里有些自定义操作，ONNX没有对应算子。解决办法是用torch.onnx.symbolic注册自定义导出规则，或者干脆改模型结构。
动态控制流：如果模型里有if语句或for循环，导出时可能出问题。ONNX要求计算图是静态的，动态控制流需要特殊处理。
精度差异：导出后推理结果和PyTorch不一致。这通常是因为某些算子的实现有细微差别。我建议导出后一定要做精度对比。

我曾经遇到过一个案例：模型里用了torch.where，导出后ONNX Runtime的结果和PyTorch差了0.01。查了半天，发现是ONNX的Where算子在处理边界情况时和PyTorch实现不同。最后我改用了torch.masked_select才解决。

2.5 验证导出的ONNX模型

导出完别急着用，先验证一下。我一般做三步验证：

结构验证：用onnx.checker.check_model检查模型结构是否合法
精度验证：用ONNX Runtime跑一遍，对比PyTorch的输出
性能验证：简单测一下推理速度，看有没有异常

代码示例：

import onnx
import onnxruntime as ort
import numpy as np

# 1. 结构验证
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("结构验证通过！")

# 2. 精度验证
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)

np.testing.assert_allclose(
    torch_output.numpy(), 
    ort_outputs[0], 
    rtol=1e-03, 
    atol=1e-05
)
print("精度验证通过！")

我的习惯：精度对比时，rtol和atol别设得太严格。1e-03和1e-05通常够用。如果模型对精度特别敏感，可以适当收紧。

2.6 小结

这一章我们聊了ONNX的由来、优势，以及PyTorch转ONNX的标准流程。核心就三点：

ONNX是模型部署的“通用语言”，解决了框架互转的痛点
导出流程很简单，但参数配置有讲究
导出后一定要验证，别跳过这步

下一章我们会深入ONNX的算子集和自定义导出，到时候再聊更高级的玩法。嗯，今天就到这儿。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录