4、ONNX实战：使用torch.onnx.export、ONNX模型验证、常见导出错误与修复

好，咱们直接进入正题。ONNX 这东西，说白了就是模型界的“通用语言”。你拿 PyTorch 训好的模型，想部署到手机上、服务器上，或者换个推理框架跑，总得有个中间格式吧？ONNX 就是干这个的。

我个人习惯，模型训完第一件事就是转 ONNX。为什么？因为 PyTorch 的动态图虽然灵活，但部署时反而是累赘。ONNX 是静态图，优化空间大，跑起来也快。今天我就带你走一遍完整的流程：导出、验证、修 bug。

核心流程： PyTorch 模型 → torch.onnx.export → ONNX 模型 → onnxruntime 验证 → 修复常见错误

PyTorch 模型 torch.onnx.export ONNX 模型 onnxruntime 验证 常见错误 → 修复 → 重新导出 ONNX 导出与验证流程

4.1 使用 torch.onnx.export 导出模型

先看最基础的导出代码。嗯，这里要注意，torch.onnx.export 其实是个“录制”过程——你给它一个输入，它跑一遍前向，然后把计算图记录下来。

import torch
import torchvision.models as models

# 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 构造一个 dummy 输入（batch_size=1, 3通道, 224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(
    model,                     # PyTorch 模型
    dummy_input,               # 示例输入
    "resnet18.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'}
    }
)
print("导出成功！")

这里有几个关键参数，我跟你聊聊我的理解：

  • opset_version：我个人习惯用 11 或 13。版本太低有些新算子不支持，版本太高有些旧推理框架又跑不了。折中选 11 最稳。
  • dynamic_axes：这个很重要。如果你部署时 batch size 不固定，一定要加上。我在项目中遇到过，没加 dynamic_axes，结果上线后发现只能一次推理一张图，那叫一个尴尬。
  • do_constant_folding：建议打开。它能提前把一些固定计算算好，比如 BN 层的参数融合，能省不少推理时间。

4.2 ONNX 模型验证

导出完了，你以为就结束了？别急。我曾经吃过这个亏——导出的 ONNX 文件看起来好好的，一跑推理结果全错。所以验证这步，绝对不能省。

验证分两步：结构验证和精度验证。

4.2.1 结构验证

用 onnx.checker 检查模型结构是否合法：

import onnx

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

# 检查模型结构
try:
    onnx.checker.check_model(onnx_model)
    print("模型结构验证通过！")
except onnx.checker.ValidationError as e:
    print(f"模型结构验证失败: {e}")

这个检查会看算子是否支持、输入输出是否匹配、节点连接是否正确。说白了就是语法检查。

4.2.2 精度验证

结构对了不代表结果对。还得用 onnxruntime 跑一遍，跟 PyTorch 的结果对比：

import onnxruntime as ort
import numpy as np

# 创建推理会话
ort_session = ort.InferenceSession("resnet18.onnx")

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# ONNX 推理
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))

# 对比结果
np.testing.assert_allclose(
    torch_output.numpy(), 
    ort_outputs[0], 
    rtol=1e-03, 
    atol=1e-05
)
print("精度验证通过！最大误差:", np.max(np.abs(torch_output.numpy() - ort_outputs[0])))

小技巧： 我一般设 rtol=1e-03, atol=1e-05。如果误差超过这个范围，说明导出过程可能有问题。但注意，有些算子（比如 softmax）在不同框架下实现略有差异，允许一点点误差是正常的。

4.3 常见导出错误与修复

这部分是我最想跟你分享的。我在部署一线摸爬滚打这么多年，遇到的坑十个手指头数不过来。下面这几个是最常见的：

错误类型 错误信息 原因 解决方案
算子不支持 RuntimeError: Unsupported operator: aten::xxx PyTorch 的某个算子 ONNX 不支持 替换为等价算子，或自定义 ONNX 算子
动态 shape 问题 RuntimeError: Input size mismatch 输入尺寸与模型期望的不一致 正确设置 dynamic_axes，或固定输入尺寸
控制流问题 RuntimeError: If/loop not supported 模型中有 if 或 for 循环 用 torch.where 替代 if，或用固定循环次数
数据类型问题 TypeError: Input must be float32 输入数据类型不匹配 确保输入是 float32，用 .float() 或 .astype() 转换

4.3.1 算子不支持

这是最头疼的。我记得有一次，模型里用了 torch.nn.functional.grid_sample，ONNX 死活不支持。后来我查了半天，发现可以用 torch.nn.functional.affine_grid 加 torch.nn.functional.grid_sample 的组合来替代，但 ONNX 只支持后者。

解决方案有两种：

  • 替换算子：找 ONNX 支持的等价实现。比如 torch.where 替代 if-else。
  • 自定义算子：如果实在找不到替代，可以自己写一个 ONNX 自定义算子。但说实话，这门槛有点高，我一般不建议新手碰。

4.3.2 动态 shape 问题

这个坑我踩过好几次。你想想看，模型训练时 batch size 是 32，导出时 dummy input 也是 32，一切正常。但部署时只传 1 张图，结果就报错了。

为什么会这样？因为 ONNX 默认把输入 shape 固定死了。解决办法就是前面说的 dynamic_axes 参数。但要注意，不是所有算子都支持动态 shape，比如 reshape 操作如果用了固定尺寸，就得改代码。

警告： 如果你的模型里有 torch.reshape 或 view 操作，并且用到了 -1 这种自动推断尺寸，导出时一定要小心。ONNX 的 reshape 算子对动态 shape 支持不太好，建议用 torch.nn.AdaptiveAvgPool2d 这类自适应算子替代。

4.3.3 控制流问题

PyTorch 的动态图允许你在 forward 里写 if-else 和 for 循环。但 ONNX 是静态图，它要求计算图是固定的。所以如果你模型里有：

def forward(self, x):
    if x.shape[0] > 10:
        return self.branch1(x)
    else:
        return self.branch2(x)

这种代码，导出时必报错。解决方案是用 torch.where 或 torch.max 这类张量操作来替代条件判断。

4.3.4 数据类型问题

这个最简单，但也最容易忽略。PyTorch 默认是 float32，但有时候你用了 float16 或者 double，导出时就会报类型不匹配。我建议统一用 float32，兼容性最好。

# 确保模型和输入都是 float32
model = model.float()
dummy_input = dummy_input.float()

4.4 实战建议

最后，给你几个我总结的实战建议：

  • 先小后大：先导出一个简单的子模块验证流程，再导出完整模型。这样出问题了容易定位。
  • 版本对齐：PyTorch、ONNX、onnxruntime 的版本要匹配。我一般用 PyTorch 1.10 + ONNX 1.11 + onnxruntime 1.10，这个组合比较稳。
  • 日志打开：导出时加 verbose=True 参数，能看到详细的导出日志，方便排查问题。
  • 备份原始模型：导出前一定保存好原始的 PyTorch 模型。万一 ONNX 出了问题，还能回退。

好了，ONNX 实战这部分就聊到这儿。记住，导出不是终点，验证才是关键。下次遇到导出报错，别慌，对照上面的表格一步步排查，基本都能解决。

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

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