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

好，咱们进入实战环节。

前面聊了 ONNX 是什么，为什么需要它。现在直接上手——怎么把 PyTorch 模型变成 ONNX？

说白了，核心就一个函数：torch.onnx.export。但别小看它，里面坑不少。我刚开始用的时候，也踩过几次雷。

3.1 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()
model.eval()

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

# 导出
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 示例输入
    "simple_model.onnx", # 输出路径
    input_names=["input"],
    output_names=["output"]
)

嗯，就这么几行。但这里有个关键点——模型必须处于 eval 模式。我见过有人忘了调 model.eval()，结果导出的 ONNX 里还带着 dropout 和 batchnorm 的训练行为，推理结果完全不对。

注意：导出前一定调用 model.eval()。否则 ONNX 里会包含训练时的随机行为。

3.2 动态轴与静态轴

你想想看，上面的例子导出的 ONNX，输入形状是固定的 (1, 10)。如果你传一个 (2, 10) 的输入，会报错。

这就是静态轴——所有维度都固定死了。

但实际部署中，我们经常需要处理不同 batch size 的输入。比如服务器端推理，一次来几个请求就拼成多大的 batch。这时候就需要动态轴。

怎么设置？看代码：

torch.onnx.export(
    model,
    dummy_input,
    "dynamic_model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},   # 第0维是动态的
        "output": {0: "batch_size"}
    }
)

这里 dynamic_axes 是一个字典。key 是张量名称，value 是另一个字典，指定哪些维度是动态的。

我个人习惯把 batch 维度设为动态，其他维度保持静态。为什么？

  • 动态轴越多，ONNX 优化越难，推理引擎可能生成次优代码
  • 有些算子不支持动态轴，导出时会报错
经验之谈：只对必要的维度开启动态。大多数场景下，只动态 batch size 就够了。序列长度、图像尺寸这些，尽量固定。

我曾经在一个 NLP 项目里，把序列长度也设成了动态轴。结果导出的 ONNX 在 TensorRT 上跑，性能比静态版本慢了 30%。后来老老实实固定了序列长度，速度就上来了。

3.3 opset_version 的选择

这个参数，很多人容易忽略。但它直接影响你的模型能不能顺利导出、能不能在目标设备上跑。

opset_version 说白了就是 ONNX 的算子版本。每个版本会新增一些算子，也可能废弃一些旧算子。

怎么选？我一般遵循这几个原则：

  1. 看目标推理引擎支持什么版本。比如 TensorRT 8.x 支持 opset 13~17，ONNX Runtime 支持更广。
  2. 选较新的稳定版本。新版本算子更丰富，但别选太新的——可能推理引擎还没跟上。
  3. 如果遇到导出失败，尝试降低 opset。有些 PyTorch 算子在高版本 opset 里映射不好。

举个例子：

torch.onnx.export(
    model,
    dummy_input,
    "model_opset15.onnx",
    opset_version=15  # 显式指定
)

如果不指定，默认是 9（PyTorch 1.7 之前）或 11（PyTorch 1.8+）。我个人建议至少用 opset 13 以上，因为支持了更多实用算子，比如 Softmax 的改进版本。

opset 版本 重要变化 推荐场景
9 基础算子集 兼容老旧推理引擎
11 新增 Scatter、ScatterND PyTorch 1.8 默认
13 Softmax 改进、新增 Gelu 推荐起点
15 更多 Transformer 相关算子 NLP 模型常用
17 最新稳定版 需要最新算子时
核心建议：优先选 opset 13 或 15。如果导出报错，再尝试降低到 11 或 9。

3.4 常见导出问题与排查

导出 ONNX 很少一次成功。我整理了几个高频问题：

  • Unsupported operator：PyTorch 里有些算子 ONNX 不支持。解决办法：用 torch.onnx.symbolic 自定义映射，或者换用支持的算子。
  • Shape inference 失败：动态轴设置不当。检查 dynamic_axes 是否覆盖了所有动态张量。
  • Constant folding 错误：模型里有无法静态计算的常量。尝试设置 do_constant_folding=False。

我曾经遇到一个很诡异的问题——导出的 ONNX 在 PyTorch 里推理结果正确，但用 ONNX Runtime 跑就错了。排查了半天，发现是 torch.onnx.export 默认开启了 training=TrainingMode.EVAL，但我的模型里有个自定义层在 eval 模式下行为异常。最后加了个 training=TrainingMode.PRESERVE 才解决。

调试技巧：导出后先用 onnx.checker.check_model 验证结构，再用 onnxruntime 跑一遍对比结果。两步下来，90% 的问题都能暴露。

3.5 小结

嗯，这一章内容不少。总结一下：

  • torch.onnx.export 是核心函数，记得先调 model.eval()
  • 动态轴只开必要的，别贪多
  • opset_version 选 13 或 15，遇到问题再降
  • 导出后一定要验证，别直接部署

下一章，咱们聊聊 ONNX 模型怎么优化和调试。到时候我会分享一些实际项目里踩过的坑，保证有用。

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