3、PyTorch模型导出ONNX：torch.onnx.export详解、动态轴与静态轴设置、输入输出命名规范、export常见报错与解决
好，咱们直接进入正题。模型训练完了，要部署到生产环境，第一步就是把 PyTorch 模型转成 ONNX。这一步看似简单，但坑是真不少。我刚开始做部署的时候，光一个导出就折腾了两天，后来才发现是动态轴没设对。
今天这一章，我把 torch.onnx.export 的用法、动态轴静态轴的区别、命名规范，还有那些常见的报错，一次性给你讲透。
3.1 torch.onnx.export 核心参数详解
先看最常用的调用方式。说白了，就一行代码：
import torch
# 假设你有一个训练好的模型
model = MyModel()
model.eval()
# 构造一个 dummy input，形状要和实际输入一致
dummy_input = torch.randn(1, 3, 224, 224)
# 导出
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 示例输入
"model.onnx", # 输出文件名
export_params=True, # 是否导出模型参数，默认True
opset_version=11, # ONNX算子集版本，我一般用11或13
do_constant_folding=True, # 常量折叠优化
input_names=["input"], # 输入张量名称
output_names=["output"], # 输出张量名称
dynamic_axes={} # 动态轴设置，后面细讲
)
这里有几个参数我要重点说一下：
- opset_version：这个决定了 ONNX 支持哪些算子。版本太低，有些 PyTorch 算子导不出来；版本太高，推理引擎可能不支持。我个人习惯用 11 或 13，兼容性最好。
- do_constant_folding：建议打开。它会把模型里一些固定的计算提前算好，比如 BN 层的参数融合，能减少推理时的计算量。
- export_params：一般保持 True。除非你只想导出网络结构，不导出权重，但这种情况很少见。
3.2 动态轴与静态轴设置
这是最容易踩坑的地方。什么叫动态轴？说白了，就是模型输入输出的某个维度是可变的。最常见的场景就是 batch size 不固定。
举个例子，你训练时 batch size 是 32，但生产环境可能一次只来 1 张图，或者一次来 100 张图。如果你导出时把 batch 维度设成静态的，那推理时只能接受 32 的倍数，这就很尴尬了。
动态轴的设置方法如下：
# 定义动态轴：batch 维度是动态的
dynamic_axes = {
"input": {0: "batch_size"}, # 第0维是batch，名字叫batch_size
"output": {0: "batch_size"} # 输出也一样
}
torch.onnx.export(
model,
dummy_input,
"model_dynamic.onnx",
input_names=["input"],
output_names=["output"],
dynamic_axes=dynamic_axes
)
这里要注意几点：
- 动态轴用字典表示，key 是张量名称，value 是维度索引到名称的映射。
- 输入和输出的动态轴要对应上。比如输入 batch 是动态的，输出 batch 也必须是动态的。
- 不是所有维度都能设成动态。有些算子对形状有固定要求，比如全连接层要求输入特征维度固定。
那什么时候用静态轴呢？如果你的生产环境输入形状完全固定，比如人脸识别模型输入永远是 112x112 的单张图，那就用静态轴。性能最好，也最省心。
3.3 输入输出命名规范
命名这事，看着小，但影响很大。ONNX 模型里的张量名称，是后续推理引擎识别输入输出的唯一标识。命名不规范，部署时各种报错。
我总结了几条经验：
- 名称要有意义：别用 "input1"、"output1" 这种。用 "image"、"text"、"logits"、"boxes" 这种一看就懂的。
- 不要用特殊字符：ONNX 名称支持字母、数字、下划线，但别用空格、连字符、中文。我在项目中见过有人用中文命名，结果 TensorRT 直接报错。
- 输入输出名称要唯一：同一个模型里，不能有两个张量叫同一个名字。
- 多输入多输出要对应：如果你的模型有多个输入，比如一个图像一个文本，命名要清晰区分。
# 多输入多输出的命名示例
torch.onnx.export(
model,
(image_input, text_input), # 多个输入用元组
"multi_io_model.onnx",
input_names=["image", "text"],
output_names=["class_logits", "bbox_coords"],
dynamic_axes={
"image": {0: "batch_size"},
"text": {0: "batch_size"},
"class_logits": {0: "batch_size"},
"bbox_coords": {0: "batch_size"}
}
)
3.4 export 常见报错与解决
这部分我直接给你列个表，都是我在实际项目中遇到过的坑：
| 报错信息 | 原因 | 解决方案 |
|---|---|---|
RuntimeError: ONNX export failed: Couldn't export operator xxx |
某个 PyTorch 算子 ONNX 不支持 | 升级 opset 版本，或者用 torch.onnx.register_custom_op 注册自定义算子 |
RuntimeError: Input 0 of node xxx has dynamic shape |
某个中间张量的形状是动态的，但 ONNX 要求静态 | 检查模型里是否有 reshape 或 gather 操作导致形状不确定，尝试固定它 |
RuntimeError: The size of tensor a (32) must match the size of tensor b (64) |
dummy_input 的形状和模型实际输入不匹配 | 检查模型 forward 函数的输入要求，确保 dummy_input 形状正确 |
UserWarning: No opset version specified. Defaulting to 9 |
没有指定 opset 版本，用了默认的 9 | 显式指定 opset_version=11 或更高 |
RuntimeError: ONNX export failed: Unsupported input format |
dummy_input 类型不对，比如传了 list 而不是 tensor | 确保 dummy_input 是 torch.Tensor 或 tuple of tensors |
重点： 遇到导出报错，先别慌。90% 的情况是以下三个原因之一：
- opset 版本太低，不支持某些算子
- 动态轴设置不对，导致形状推导失败
- 模型里有 if-else 分支，导致导出时走了未定义的路径
我建议你导出前先用 torch.jit.trace 试一下，它能提前暴露很多问题。
还有一个常见问题：导出成功了，但推理结果不对。这通常是因为模型里有 BatchNorm 或 Dropout 层，导出时没有切换到 eval 模式。记住，导出前一定要调用 model.eval()，不然 BN 层的 running_mean 和 running_var 不会固定，导出的权重就是错的。
嗯，这一章的内容就这些。总结一下：torch.onnx.export 的核心就三个点——参数要熟、动态轴要慎、命名要规范。把这三点吃透了，导出 ONNX 基本不会出大问题。