3、ONNX模型导出：PyTorch模型转ONNX、ONNX算子兼容性检查、动态轴与静态轴设置、ONNX Runtime验证

说实话，ONNX 这个中间格式，我一开始是有点抗拒的。总觉得多此一举——PyTorch 直接推理不香吗？直到我在一个项目里需要把模型部署到嵌入式设备上，才发现 PyTorch 的 runtime 在那上面根本跑不动。嗯，ONNX 就成了绕不开的一环。

这一章，咱们就聊聊怎么把 PyTorch 模型转成 ONNX。我会把我在实际踩坑中总结的经验都抖出来。

3.1 PyTorch 模型转 ONNX：基础操作

PyTorch 官方提供了 torch.onnx.export 这个接口。说白了，就是把你训练好的模型“拍个快照”，然后存成 ONNX 格式。

我习惯这么写：

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = MyModel()
model.eval()  # 别忘了切到 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=17,    # ONNX opset 版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=None    # 先设为 None，后面再讲动态轴
)

这里有个细节：opset_version 我建议用 17 或更高。为什么？因为低版本 opset 不支持一些新算子，比如 LayerNorm 的某些变体。我在一个 BERT 模型上吃过亏，opset=11 死活导不出，升到 17 就过了。

小提示： 导出前一定要调用 model.eval()。否则 BatchNorm 和 Dropout 的行为会不一样，导出的 ONNX 在推理时可能结果不对。

3.2 ONNX 算子兼容性检查

模型导出成功，不代表万事大吉。你想想看，ONNX 支持的算子是有限的。PyTorch 里有些骚操作，比如 torch.einsum、torch.where 的某些用法，ONNX 可能不认识。

我一般用两种方法检查：

  1. 直接看导出日志：torch.onnx.export 会在控制台打印警告，告诉你哪些算子被 fallback 了。
  2. 用 onnx.checker 验证：
import onnx

# 加载 ONNX 模型
onnx_model = onnx.load("model.onnx")

# 检查模型结构是否合法
onnx.checker.check_model(onnx_model)

print("模型检查通过！")

如果 checker 报错，别慌。常见的问题有：

  • 算子不支持：比如 torch.nn.functional.scaled_dot_product_attention 在低版本 opset 里没有。解决办法是升级 opset，或者手动拆成基础算子。
  • 输入输出形状不匹配：检查一下你的 dummy_input 形状对不对。
  • 动态轴设置冲突：这个我们马上讲。
注意： 我曾经遇到过一个坑——模型导出时没报错，checker 也通过了，但推理结果完全不对。后来发现是 torch.where 被 ONNX 转换成了 NonZero 算子，语义变了。所以，光检查还不够，还得用 ONNX Runtime 跑一遍验证。

3.3 动态轴与静态轴设置

静态轴，就是固定死的形状。比如你的模型只接受 batch_size=1 的输入。动态轴，就是允许输入形状变化，比如 batch_size 可以是 1、2、4、8……

我个人习惯：能静态就静态，需要动态才动态。因为动态轴会引入额外的 shape 计算，推理速度会慢一点。

设置动态轴的方法：

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
)

这里 {0: 'batch_size'} 表示第 0 维是动态的，名字叫 batch_size。你可以在推理时传入不同 batch_size 的输入。

但要注意：不是所有算子都支持动态轴。比如 Reshape 如果依赖某个动态维度的值，就可能出问题。我遇到过 Gather 算子在动态轴下索引越界的情况，排查了半天。

经验之谈： 如果你不确定哪些维度需要动态，可以先全部设成静态。等部署时发现 batch_size 不够用，再回来改。别一开始就搞得太复杂。

3.4 ONNX Runtime 验证

导出 ONNX 只是第一步。你得用 ONNX Runtime 跑一遍，确保结果和 PyTorch 一致。

我的验证流程：

import onnxruntime as ort
import numpy as np

# 创建 ONNX Runtime session
ort_session = ort.InferenceSession("model.onnx")

# 准备输入数据（用和 PyTorch 一样的随机种子）
np.random.seed(42)
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outputs = ort_session.run(None, ort_inputs)

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

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

if diff < 1e-5:
    print("验证通过！")
else:
    print("验证失败，请检查模型导出过程。")

这里有个坑：ONNX Runtime 默认用 CPU 推理。如果你模型里有 torch.cuda.FloatTensor，导出时记得把 dummy_input 也放到 CPU 上。否则导出的 ONNX 会包含 CUDA 相关的算子，ONNX Runtime 跑不了。

小技巧： 我一般会在验证时多测几组不同输入，包括边界情况（比如全零输入、全一输入）。这样能发现一些隐藏的 bug。

3.5 知识体系总览

这一章的内容比较多，我画了个图帮你理清思路：

ONNX模型导出知识体系 ONNX模型导出 PyTorch转ONNX torch.onnx.export 算子兼容性检查 onnx.checker.check_model 动态轴与静态轴 dynamic_axes参数 ONNX Runtime验证 onnxruntime.InferenceSession opset_version 推荐≥17 do_constant_folding 常量折叠优化 检查日志警告 算子fallback 常见问题 算子不支持/形状不匹配 动态维度命名 如batch_size 性能权衡 静态更快 数值对比 误差<1e-5 多组测试 边界情况 导出 → 检查 → 验证 → 部署

这张图把整个流程串起来了。从 PyTorch 导出开始，到算子兼容性检查，再到动态轴设置，最后用 ONNX Runtime 验证。每一步都有对应的工具和方法。

嗯，这一章的内容就这些。记住：导出不是终点，验证才是。我见过太多人导出成功就以为完事了，结果部署时发现推理结果不对，又回头排查。别走这个弯路。

公众号:蓝海资料掘金营，微信deep3321

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