第4章 ONNX模型导出（PyTorch）：torch.onnx.export深度拆解
把PyTorch模型转成ONNX，这事我干过不下几百次了。说实话，最开始踩的坑能绕地球半圈。torch.onnx.export这个接口，参数看着不多，但每一个都可能让你的模型导出失败，或者跑起来结果不对。今天我把这些参数掰开了讲，顺便说说我踩过的那些坑。
4.1 torch.onnx.export核心参数
先看最常用的调用方式，我直接贴代码：
torch.onnx.export(
model, # 要导出的模型
args, # 模型的输入张量（或元组）
"model.onnx", # 输出文件路径
export_params=True, # 是否导出可训练参数
opset_version=17, # ONNX算子集版本
do_constant_folding=True, # 是否做常量折叠优化
input_names=["input"], # 输入节点名称
output_names=["output"], # 输出节点名称
dynamic_axes={}, # 动态轴设置
verbose=False, # 是否打印导出日志
)我习惯把每个参数都显式写出来，虽然有些有默认值，但显式写清楚，团队其他人接手也不会懵。下面我挑几个容易出问题的参数细说。
| 参数名 | 默认值 | 我的建议 |
|---|---|---|
| opset_version | 取决于PyTorch版本 | 尽量用最高稳定版（目前17），低版本可能缺算子 |
| do_constant_folding | True | 除非你明确知道要关，否则别动 |
| export_params | True | 推理模型一般保留 |
| verbose | False | 第一次导出建议True，方便debug |
4.2 opset_version到底选多少？
我记得有一次项目急着上线，我选了最新的opset=18，结果TensorRT根本不吃，提示算子不支持。后来灰溜溜降到了17。我个人经验：用opset=17最稳妥。因为TensorRT的主流版本对17的支持最成熟。如果你不确定，先查一下目标推理框架的官方文档支持到哪个版本。
4.3 dynamic_axes：动态轴的正确姿势
什么是动态轴？说白了就是允许模型在推理时接受不同大小的输入。比如你训练时用了batch=1，但部署时想支持batch=4、8、16，那就得设动态batch。
看个例子：
dynamic_axes = {
"input": {0: "batch_size", 2: "height", 3: "width"},
"output": {0: "batch_size"}
}
torch.onnx.export(
model,
dummy_input,
"model_dynamic.onnx",
input_names=["input"],
output_names=["output"],
dynamic_axes=dynamic_axes,
opset_version=17,
)注意这里的key是输入/输出名称，必须和input_names/output_names里的一致。字典的value是轴索引: 轴名称的映射。轴名称你自己随便起，但要有可读性。
有人问：动态轴设多了会不会影响性能？会的。你想想看，每个动态轴都会让TensorRT多生成一份中间表示，推理开销略大。我的原则是：只设必须动态的维度。比如你的模型只接受固定大小的图像，那就别设height/width的动态。
4.4 算子支持列表与替代方案
ONNX定义了一大堆算子，但PyTorch有些操作并没有直接映射。我遇到过最典型的几个：
| PyTorch操作 | ONNX支持情况 | 替代方案 |
|---|---|---|
| torch.einsum | 部分支持（opset>=16） | 改用torch.matmul/手动reshape |
| torch.where | opset>=9 | 写成mask * x + (1-mask) * y |
| torch.topk | opset>=11 | 保持argmax/手动排序 |
| F.interpolate | 支持但mode有限 | 使用'nearest'或'bilinear'，避免'area' |
| torch.split | opset>=8，但动态split容易出问题 | 改用torch.chunk或手动切片 |
我踩过一个最深的坑：torch.cdist（计算点对距离），在opset=17下导出的模型居然不报错，但结果全是0。后来我查了ONNX算子列表，发现cdist压根没实现。最后我用torch.nn.PairwiseDistance替代，才搞定。
4.5 导出的常见报错与处理
- RuntimeError: Unsupported operator: 直接告诉你哪个算子不支持。去PyTorch的ONNX支持列表查一下，或者用torch.onnx.register_custom_op_symbolic手动注册。
- TypeError: 0D tensor expected: 通常是因为某个标量操作返回了0维张量。加个
.unsqueeze(0)就行。 - ValueError: Could not export operator: 如果你用了控制流（if/for），得用torch.jit.script或torch.jit.trace来处理。
我一般在模型代码里加一个export_mode的flag，导出时启用，把不支持的算子替换成支持的版本。比如：
def forward(self, x):
if self.export_mode:
return torch.matmul(x, self.weight)
else:
return torch.einsum('...ij,...jk->...ik', x, self.weight)这样训练和导出用不同路径，互不影响。
4.6 一张图理清ONNX导出流程
我画了张流程图，把从PyTorch模型到ONNX文件的关键步骤和决策点串起来。你看一眼就懂了。这张图必须和本章强相关：核心就是带你走一遍导出时的参数设置、动态轴、算子检查、最终生成ONNX的路径。
4.7 提几个落地建议
- 先跑一个小模型：用单层网络先测试导出，确认基础流程通了，再用完整模型。
- 验证精度：导出后一定要用ONNX Runtime或TensorRT跑一遍推理，对比PyTorch原输出，差值超过1e-5就要查原因。
- 保持代码整洁：用config字典管理export参数，方便团队复现和修改。
嗯，关于torch.onnx.export的知识点就这些了。说来说去，核心就三件事：设对参数、管好动态轴、搞定算子替代。你在实际项目中遇到问题，优先检查这三块，十有八九能解决。