4、ONNX与PyTorch的关系：PyTorch原生导出ONNX的机制
聊到ONNX，就绕不开PyTorch。这两者的关系，说白了就是「亲兄弟」。
PyTorch现在是AI框架里的顶流，而ONNX是它官方支持的模型交换格式。我刚开始接触ONNX时，总觉得这玩意儿就是个「中间人」，后来才发现——它其实是PyTorch生态里不可或缺的一环。
4.1 PyTorch为什么选择ONNX？
你想想看，PyTorch模型训练好了，总不能永远待在Python环境里吧？
生产环境可能是C++、Java，甚至是移动端或嵌入式设备。PyTorch自己有一套TorchScript，但说实话，它的生态覆盖面和ONNX比，还是有差距的。
我个人习惯是：训练用PyTorch，部署用ONNX。这样既享受了PyTorch的灵活，又拿到了ONNX的跨平台能力。
4.2 原生导出机制：torch.onnx.export
PyTorch从1.2版本开始，就内置了ONNX导出功能。你不需要装任何第三方库，直接调 torch.onnx.export 就行。
来看一个最简单的例子：
import torch
import torch.nn as nn
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = SimpleModel()
dummy_input = torch.randn(1, 10)
torch.onnx.export(
model, # 模型
dummy_input, # 示例输入
"simple_model.onnx", # 输出路径
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}}
)
这段代码我写过不下百次。嗯，这里要注意几个关键点：
- dummy_input 必须和实际输入的形状一致，否则导出会报错
- dynamic_axes 用来指定动态维度，比如batch_size可以变化
- input_names/output_names 给张量起个名字，方便后续推理时引用
4.3 导出背后的原理：Tracing vs Scripting
PyTorch导出ONNX，底层用的是 Tracing（追踪） 机制。
什么意思呢？就是给模型一个假输入，让它跑一遍，然后记录下所有操作。说白了，就是「看你怎么算，我就怎么记」。
但这里有个坑——控制流。如果模型里有if-else或者循环，Tracing只会记录当前分支的路径。比如：
class MyModel(nn.Module):
def forward(self, x):
if x.sum() > 0:
return self.fc1(x)
else:
return self.fc2(x)
如果你给的dummy_input是正数，导出的ONNX里就只有fc1那条路。部署时遇到负数输入，模型就傻了。
解决办法有两个：
- 用
torch.jit.script先把模型转成TorchScript，再导出ONNX - 或者手动保证dummy_input覆盖所有分支
4.4 动态轴与静态轴
ONNX默认是静态图，所有维度都是固定的。但实际部署时，输入batch_size经常变化。
PyTorch的 dynamic_axes 参数就是干这个的。我建议你养成习惯，只要不确定输入大小，就加上这个参数。
| 参数 | 作用 | 示例 |
|---|---|---|
| dynamic_axes | 指定哪些维度可以动态变化 | {'input': {0: 'batch_size'}} |
| opset_version | 指定ONNX算子集版本 | opset_version=17 |
| do_constant_folding | 是否做常量折叠优化 | do_constant_folding=True |
这里有个小技巧：opset_version 尽量选新版本。新版本支持更多算子，兼容性更好。我一般用17或18。
4.5 算子映射：PyTorch → ONNX
PyTorch的每个操作，都要映射到ONNX的算子。大部分常见算子（Conv、ReLU、BatchNorm）都有对应关系。
但总有些「漏网之鱼」。比如 torch.einsum，ONNX早期版本就不支持。
torch.onnx.register_custom_op 注册自定义算子。或者干脆换个实现方式，比如把einsum拆成矩阵乘法。
PyTorch官方维护了一个算子映射表，我建议你导出前先查一下。特别是用了比较新的API时，比如 torch.nn.functional.scaled_dot_product_attention，在旧版ONNX里可能找不到对应算子。
4.6 验证导出结果
导出ONNX后，千万别直接拿去用。先验证一下：
import onnx
import onnxruntime as ort
import numpy as np
# 检查模型结构
onnx_model = onnx.load("simple_model.onnx")
onnx.checker.check_model(onnx_model)
# 对比推理结果
ort_session = ort.InferenceSession("simple_model.onnx")
input_data = np.random.randn(1, 10).astype(np.float32)
ort_output = ort_session.run(None, {'input': input_data})
torch_output = model(torch.from_numpy(input_data)).detach().numpy()
# 检查误差
np.testing.assert_allclose(ort_output[0], torch_output, rtol=1e-3, atol=1e-5)
我个人的习惯是：误差控制在1e-5以内。如果偏差太大，八成是算子映射出了问题。
4.7 常见问题与避坑
- 动态图问题： PyTorch是动态图，ONNX是静态图。导出时会把动态图「拍平」成静态图，有些动态行为会丢失。
- 数据类型： ONNX默认用float32。如果你的模型用了float16或double，导出时要注意类型转换。
- BatchNorm和Dropout： 导出前记得调用
model.eval()，否则BN和Dropout的行为会不一致。
model.eval()，结果ONNX模型里还带着Dropout，推理结果每次都不一样。排查了半天才发现是训练模式和推理模式没切换。
4.8 总结一下
PyTorch和ONNX的关系，就像「作者」和「出版社」。PyTorch负责创作，ONNX负责分发。
导出机制本身不复杂，核心就三步：
- 准备模型和示例输入
- 调用
torch.onnx.export - 验证导出结果
但细节决定成败。动态轴、算子映射、验证精度——这些环节一个都不能省。
好了，这一章就到这里。记住一句话：导出不是终点，验证才是关键。