📚 课程目录第3章 PyTorch模型导出ONNX

3、PyTorch模型导出ONNX：torch.onnx.export详解、动态轴与静态轴、opset版本选择

好，咱们进入实战环节的第一个硬骨头——把PyTorch模型转成ONNX。

说实话，很多同学在这个阶段就开始踩坑了。明明PyTorch里跑得好好的模型，一导出就报错，或者导出来推理结果不对。我当年第一次做这个转换时，也折腾了两天才搞明白问题出在哪。

这一节，咱们就把torch.onnx.export这个函数彻底讲透。包括它的每个参数怎么用、动态轴和静态轴的区别、以及opset版本该怎么选。

3.1 torch.onnx.export 核心参数详解

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

import torch
import torch.onnx

# 假设我们有一个训练好的模型
model = torch.load('my_model.pth')
model.eval()

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

# 导出 ONNX
torch.onnx.export(
    model,                    # 要导出的模型
    dummy_input,              # 示例输入
    'model.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'}
    }
)

这段代码看着简单，但每个参数背后都有讲究。我一个个说。

3.2 参数逐个拆解

参数名	作用	我的建议
`model`	要导出的PyTorch模型	导出前务必调用`model.eval()`，关闭dropout和batch norm的训练行为
`args`	示例输入（tuple或tensor）	形状和数据类型要与实际推理时一致，否则导出后可能shape mismatch
`export_params`	是否导出模型权重	默认True。如果只想导出网络结构，设为False
`opset_version`	ONNX算子集版本	我一般用11或13，兼容性和功能平衡得比较好
`do_constant_folding`	是否做常量折叠优化	建议开启，能减小模型体积、提升推理速度
`input_names`	输入节点名称	命名要有意义，方便后续调试
`output_names`	输出节点名称	同上
`dynamic_axes`	定义哪些维度是动态的	端侧部署时几乎必用，后面细讲

我的小技巧： 导出前先用torch.onnx.export的verbose=True模式跑一次，能看到完整的计算图结构，方便排查问题。

3.3 动态轴 vs 静态轴

这是最容易搞混的概念。我直接说人话：

静态轴：导出时输入输出的形状就固定死了。比如[1, 3, 224, 224]，推理时也只能用这个形状。
动态轴：允许某些维度在推理时变化。最常见的就是batch size维度。

举个例子。你训练时batch size是32，但部署到嵌入式设备上，可能一次只推理1张图。如果导出时用了静态轴，那ONNX模型就只接受batch size=32的输入。这就很尴尬了。

所以，我建议在端侧部署时，至少把batch size维度设为动态。代码里这样写：

dynamic_axes = {
    'input': {0: 'batch_size'},   # 第0维是batch size
    'output': {0: 'batch_size'}
}

如果你模型里有可变长度的输入（比如NLP任务里的sequence length），也可以把那个维度设成动态：

dynamic_axes = {
    'input': {0: 'batch_size', 1: 'seq_len'},
    'output': {0: 'batch_size'}
}

注意： 动态轴不是越多越好。每个动态轴都会增加推理引擎的调度开销。我见过有人把4个维度全设成动态，结果推理速度掉了30%。够用就好。

3.4 opset版本怎么选？

opset版本，说白了就是ONNX算子集的版本号。不同版本支持的算子种类和功能不一样。

我整理了一个常用版本的对比：

opset版本	关键特性	适用场景
9	基础算子完备，兼容性最好	老旧推理引擎、ARM Cortex-M等资源受限设备
11	支持更多高级算子，如`ScatterND`	大部分嵌入式Linux设备，我的首选
13	支持`Softmax`的axis参数等改进	需要用到较新算子的模型
15+	支持FP16、BF16等精度控制	有量化需求的场景

我个人习惯用opset 11。为什么？

它支持的算子足够多，覆盖了90%以上的常见模型
兼容性好，主流的推理引擎（ONNX Runtime、TensorRT、OpenVINO）都支持得很好
踩坑经验少，社区资料丰富

如果你用了比较新的模型结构（比如Transformer的一些变体），可能需要opset 13或更高。但记住一个原则：能用低版本就别用高版本。版本越高，对推理引擎的要求也越高。

    我曾经踩过的坑： 有一次用opset 15导出一个模型，在ONNX Runtime上跑得好好的，但部署到某个嵌入式平台的NPU上时，发现它不支持opset 15的某个算子。最后只能降级到opset 11，手动替换了几个算子才搞定。所以，先确认目标平台的opset支持范围，再选版本。

3.5 导出流程全景图

下面这张图把整个导出流程串起来了，你可以对照着看：

3.6 导出后的验证

模型导出后，别急着部署。先做两件事：

用ONNX Runtime验证精度

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
sess = ort.InferenceSession('model.onnx')

# 构造输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
outputs = sess.run(None, {'input': input_data})

# 对比PyTorch结果
with torch.no_grad():
    torch_output = model(torch.from_numpy(input_data)).numpy()

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

如果误差在1e-5以内，基本没问题。如果误差很大，可能是某些算子转换出了问题。

用Netron可视化检查网络结构

Netron是一个ONNX可视化工具，能直观看到网络结构。我每次导出后都会打开看一眼，确认输入输出节点、shape信息都正确。

我的习惯： 在导出脚本里加一个自动验证步骤。导出完成后，自动用ONNX Runtime跑一次推理，对比PyTorch结果。这样能第一时间发现问题，不用等到部署时再排查。

3.7 常见问题与避坑

算子不支持：如果导出时报错"Unsupported operator"，说明PyTorch里用了ONNX不支持的算子。解决办法是替换成等价算子，或者用torch.onnx.register_custom_op注册自定义算子。
动态轴导致性能下降：动态轴越多，推理引擎的优化空间越小。我建议只把必要的维度设成动态，其他维度固定。
opset版本不匹配：导出时用的opset版本，推理引擎必须支持。否则会报"Unsupported opset version"错误。

嗯，这一节的内容就这些。核心就是三个东西：torch.onnx.export的参数、动态轴与静态轴的选择、opset版本的权衡。把这些搞明白，导出ONNX这一步就不会出大问题。