4、PyTorch模型导出ONNX：torch.onnx.export详解、动态轴与静态轴设置、常见导出错误与修复

好，咱们进入实战环节。前面几章我们把TensorRT的底层原理、模型优化策略都聊了一遍，但有个关键问题一直没解决——模型怎么从PyTorch跑到TensorRT里去？

答案就是ONNX。它像个中间人，负责把PyTorch的模型描述翻译成TensorRT能理解的语言。但说实话，这个翻译过程坑不少。我刚开始做部署时，光导出ONNX就卡了整整两天。今天我把这些经验全盘托出，你照着做，能少走很多弯路。

4.1 torch.onnx.export 核心参数解析

先看最常用的导出方式。PyTorch官方提供了torch.onnx.export这个接口，用法其实不复杂：

import torch
import torch.onnx

# 假设我们有一个训练好的模型
model = MyModel()
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算子集版本
    do_constant_folding=True,  # 是否做常量折叠优化
    input_names=['input'],     # 输入张量名称
    output_names=['output'],   # 输出张量名称
    dynamic_axes={       # 动态轴设置
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

嗯，这里要注意几个关键参数。我个人习惯把opset_version设成11或更高，因为低版本有些算子不支持。但也不是越高越好——我遇到过opset 13导出的模型在旧版TensorRT上跑不了的情况。所以我的建议是：先试11，不行再往上调。

do_constant_folding这个参数，说白了就是让ONNX在导出时把能算的常量先算好。比如模型里有torch.full([3], 1.0)这种操作，导出时直接变成常量，推理时省一步计算。我一般默认开着，除非遇到奇怪的导出错误。

4.2 动态轴与静态轴：你该怎么选？

这是最容易踩坑的地方。先搞清楚概念：

  • 静态轴：输入输出的形状在导出时就固定死了。比如batch size永远是1，图片永远是224x224。
  • 动态轴：某些维度可以在推理时变化。比如batch size可以是1、4、8、16任意值。

为什么要区分？因为TensorRT对静态形状的优化更激进。你想想看，如果形状固定，TensorRT可以做很多内存布局的预计算，推理速度能快10%-20%。

但实际场景中，我们经常需要动态batch。比如线上服务同时来了多个请求，你总不能一个一个推理吧？这时候就需要动态轴。

看个具体例子：

# 静态导出：batch size固定为1
torch.onnx.export(
    model,
    torch.randn(1, 3, 224, 224),
    "model_static.onnx",
    input_names=['input'],
    output_names=['output']
)

# 动态导出：batch size可变
torch.onnx.export(
    model,
    torch.randn(1, 3, 224, 224),
    "model_dynamic.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

注意看区别。静态导出时，ONNX文件里写死了输入形状是[1,3,224,224]。动态导出时，第0维被标记为可变的batch_size。

我在项目中遇到过一个问题：某个模型用了torch.view操作，动态batch导出后，TensorRT报错说形状不匹配。后来排查发现，view操作在动态形状下会生成Reshape节点，而TensorRT对Reshape的处理有坑。解决方案是把view改成reshape，或者用-1推断维度。

我的建议：能静态就别动态。如果必须动态，尽量只让batch size可变，其他维度（宽高、通道数）保持静态。这样TensorRT的优化效果最好。

4.3 常见导出错误与修复

这部分我花了不少时间总结。你遇到的大部分导出问题，下面都能找到答案。

错误1：算子不支持

错误信息类似：RuntimeError: ONNX export failed: Couldn't export operator xxx

为什么会这样？因为PyTorch有些操作ONNX标准里没有。比如torch.einsum、torch.topk的某些模式、自定义的C++扩展等。

解决方案：

  • 升级opset版本（试试13、14、15）
  • 用PyTorch官方提供的torch.onnx.symbolic注册自定义导出规则
  • 实在不行，在导出前把模型里的特殊操作替换成ONNX支持的等价操作

我曾经遇到一个模型用了torch.meshgrid，opset 11不支持，升级到13就好了。所以遇到算子问题，先试试升级opset。

错误2：动态形状导致图结构错误

错误信息：RuntimeError: The size of tensor a (X) must match the size of tensor b (Y) at non-singleton dimension Z

这个错误很隐蔽。我排查过整整一天才发现问题：模型里有个torch.cat操作，拼接的两个张量在动态batch下形状推导出了问题。

修复方法：

  • 检查所有涉及形状的操作（cat、stack、view、expand等）
  • 用torch.onnx.export的verbose=True参数打印详细日志
  • 使用torch.onnx.export的check_trace=True做运行时检查

错误3：控制流导致导出失败

错误信息：RuntimeError: ONNX export failed: Could not export a Python object

PyTorch模型里如果有if语句、for循环这些Python控制流，ONNX导出时会直接报错。因为ONNX是静态图，不支持动态控制流。

解决方案：

  • 把控制流改成torch.where或torch.max等张量操作
  • 如果必须用控制流，考虑用torch.jit.script先转成TorchScript，再导出ONNX
注意：用TorchScript导出ONNX虽然能解决控制流问题，但可能会引入新的兼容性问题。我建议只在万不得已时用这个方案。

错误4：输入输出命名冲突

这个错误比较低级，但容易忽略。如果你在dynamic_axes里用了'input'这个名字，但模型里恰好有个张量也叫'input'，就会冲突。

解决方案很简单：给输入输出起个独一无二的名字，比如'input_0'、'output_0'，或者用模型里实际张量的名字。

4.4 导出后的验证与调试

导出ONNX只是第一步。我强烈建议你导出后做两件事：

  1. 用ONNX Runtime验证：确保导出的ONNX能正确推理
  2. 对比精度：用同样的输入，对比PyTorch和ONNX的输出差异

验证代码示例：

import onnx
import onnxruntime as ort
import numpy as np

# 1. 检查ONNX模型结构
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过")

# 2. 用ONNX Runtime推理
ort_session = ort.InferenceSession("model.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: np.random.randn(1, 3, 224, 224).astype(np.float32)}
ort_outputs = ort_session.run(None, ort_inputs)

# 3. 对比PyTorch结果
with torch.no_grad():
    torch_output = model(torch.from_numpy(ort_inputs['input']))

# 计算最大误差
max_diff = np.max(np.abs(ort_outputs[0] - torch_output.numpy()))
print(f"PyTorch与ONNX最大误差：{max_diff:.6f}")

如果最大误差超过1e-5，说明导出过程有问题。我遇到过因为do_constant_folding导致精度下降的情况，关掉就好了。

4.5 实战经验总结

最后，分享几个我踩过的坑：

  • BatchNorm层要冻结：导出前一定要model.eval()，否则BatchNorm的running mean和running var会乱掉
  • Dropout层要关闭：同样，model.eval()会关闭Dropout，导出时不会包含这些训练专用操作
  • dummy input的形状要准确：不要用torch.randn(1, 3, 224, 224)去测试输入为(1, 3, 256, 256)的模型，会导出错误的形状信息
  • 多输出模型要命名清晰：如果模型有多个输出，给每个输出起个有意义的名字，方便后续在TensorRT里绑定

嗯，导出ONNX这一步看似简单，但细节决定成败。我见过太多人卡在这一步，然后怀疑人生。其实只要按照上面的步骤来，大部分问题都能解决。下一章我们聊聊怎么把ONNX模型转换成TensorRT的engine，那才是真正加速的开始。

← 上一章 📖 返回目录 下一章 →