4、ONNX基础：什么是ONNX？PyTorch模型转ONNX、ONNX算子兼容性检查

好，咱们进入第四讲。前面几章我们把目标检测的模型选型、训练流程都捋了一遍。现在模型训好了，权重文件躺在硬盘里，接下来干嘛？部署啊！

但部署有个很现实的问题：你训练用的框架是PyTorch，可生产环境可能是TensorRT、OpenVINO，或者干脆就是个手机端。总不能每换一个平台就重写一遍模型吧？

这时候，ONNX就登场了。

4.1 什么是ONNX？

ONNX，全称Open Neural Network Exchange，开放神经网络交换格式。说白了，它就是深度学习模型界的“通用语言”。

你想想看，PyTorch有它自己的模型格式（.pt/.pth），TensorFlow有它的SavedModel，Caffe有它的.prototxt。这些格式互不兼容，就像各国说各地方言。ONNX就是那个“普通话”——不管你用什么框架训练，导出成ONNX之后，任何支持ONNX的推理引擎都能读懂。

ONNX的核心价值就两点：

  • 互操作性：一次导出，到处运行。从PyTorch到TensorRT，从TensorFlow到ONNX Runtime，无缝切换。
  • 推理优化：ONNX本身是一个计算图描述格式。很多推理引擎（比如ONNX Runtime、TensorRT）会针对ONNX图做算子融合、内存优化，推理速度往往比原始框架快不少。

重要概念：ONNX不是推理引擎，它是模型描述格式。

你导出ONNX文件后，还需要用ONNX Runtime或者转成TensorRT引擎才能跑推理。别搞混了。

我个人习惯把ONNX看作“中间表示层”。训练用PyTorch，部署用ONNX，各司其职。我在项目中遇到过不少团队，直接在PyTorch里做推理部署，结果线上环境装了一堆Python依赖，模型加载慢、推理延迟高。后来统一走ONNX + ONNX Runtime，问题全解决了。

4.2 PyTorch模型转ONNX

好，理论说完了，咱们直接上手。把PyTorch模型转成ONNX，核心就一个函数：torch.onnx.export()。

先看一个最简单的例子。假设我们有一个训练好的分类模型：

import torch
import torchvision.models as models

# 1. 加载模型，设为eval模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,                    # 模型
    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("ONNX导出成功！")

这段代码看着简单，但有几个坑我得提醒你：

  • 必须调用model.eval()：不设成eval模式，BatchNorm和Dropout的行为会乱掉。我见过有人忘了这步，导出的ONNX推理结果跟PyTorch对不上，排查了半天。
  • dummy_input的shape要跟实际输入一致：尤其是目标检测模型，输入尺寸通常是固定的（比如640x640）。如果你训练时用了动态尺寸，那就要设置dynamic_axes。
  • opset_version别选太新：我一般用opset=11或12。太新的版本（比如18、19）有些推理引擎还没完全支持，导出容易报错。

小技巧：如果你不确定用哪个opset版本，可以先试试11。这是最广泛支持的版本，兼容性最好。

对于目标检测模型，比如YOLOv5或YOLOv8，导出时还要注意后处理部分。PyTorch版本的YOLO通常包含NMS（非极大值抑制）操作，但ONNX标准算子集里没有NMS算子。怎么办？

两种方案：

  1. 导出时不包含NMS：只导出Backbone + Neck + Head部分，输出原始的检测框和置信度。后处理在推理端用Python或C++实现。
  2. 使用ONNX自定义算子：把NMS封装成自定义算子，导出时带上。但这样推理端也要注册对应的算子实现，比较麻烦。

我个人建议用方案一。原因很简单：保持ONNX的纯净性。后处理代码写起来也不复杂，而且方便调试。

4.3 ONNX算子兼容性检查

模型导出成功了，是不是就万事大吉了？

别急。你导出的ONNX文件，拿到目标推理引擎上，可能跑不起来。为什么？算子不兼容。

举个例子：PyTorch里有个torch.nn.functional.grid_sample，这个操作在ONNX的opset 11里是不支持的。如果你模型里用了它，导出时要么报错，要么自动拆成多个基础算子。但拆完之后，推理引擎能不能正确执行，又是另一回事。

所以，导出之后一定要做兼容性检查。我常用的工具有两个：

4.3.1 使用ONNX Runtime检查

最简单的方法，直接用ONNX Runtime加载模型，跑一次推理：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
session = ort.InferenceSession("resnet18.onnx")

# 获取输入输出信息
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 构造输入数据
dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 跑推理
outputs = session.run([output_name], {input_name: dummy_input})

print("ONNX Runtime推理成功！")
print("输出shape:", outputs[0].shape)

如果这一步没报错，说明你的ONNX模型至少能在ONNX Runtime上跑通。但注意，这只能验证ONNX Runtime的兼容性，换到TensorRT或者OpenVINO上可能还有问题。

4.3.2 使用onnx.checker检查模型结构

ONNX官方提供了一个检查工具，可以验证模型图结构的合法性：

import onnx

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

# 检查模型结构
onnx.checker.check_model(model)

print("ONNX模型结构检查通过！")

这个检查会验证：节点连接是否正确、输入输出是否匹配、算子版本是否合法等。如果这里报错，说明你的ONNX文件本身就有问题，需要重新导出。

注意：onnx.checker只检查结构合法性，不检查数值精度。也就是说，模型结构没问题，但推理结果可能跟PyTorch不一致。这种情况我遇到过好几次，尤其是量化模型或者用了自定义算子的情况。

4.3.3 数值精度对比

这是最严格的一步。把PyTorch的输出和ONNX Runtime的输出做对比：

import torch
import onnxruntime
import numpy as np

# PyTorch推理
model.eval()
with torch.no_grad():
    pt_output = model(torch.from_numpy(dummy_input)).numpy()

# ONNX Runtime推理
session = onnxruntime.InferenceSession("resnet18.onnx")
ort_output = session.run(None, {input_name: dummy_input})[0]

# 对比
diff = np.abs(pt_output - ort_output)
print("最大误差:", diff.max())
print("平均误差:", diff.mean())

一般来说，误差在1e-5以内都是正常的。如果误差很大，说明导出过程中某些算子被替换了，或者精度丢失了。这时候就要检查模型里有没有不支持的算子，或者考虑换一个opset版本。

4.4 常见算子兼容性问题

我在项目中踩过的坑，列出来给你参考：

PyTorch算子 ONNX支持情况 建议方案
torch.nn.functional.grid_sample opset 16+ 才支持 升级opset，或用自定义算子
torch.nn.functional.interpolate(mode='bilinear') opset 11+ 支持 一般没问题，注意align_corners参数
torch.split / torch.chunk opset 11+ 支持 注意split sizes参数必须是常量
torch.nonzero opset 9+ 支持 输出shape是动态的，注意dynamic_axes设置
torch.sort / torch.topk opset 11+ 支持 NMS中常用，一般没问题

避坑指南：我曾经把一个用了torch.meshgrid的模型导出ONNX，结果在opset 11下报错。查了半天才发现，meshgrid在opset 11里只支持2个输入，而我用了3个。后来升级到opset 13才解决。所以，遇到算子问题，先查ONNX官方文档，看看你用的opset版本到底支持哪些算子。

4.5 总结

这一章我们聊了ONNX的基础知识。说白了，ONNX就是模型部署的“中间人”。它帮你把PyTorch模型转换成通用格式，然后你可以在任何支持ONNX的推理引擎上跑。

核心要点就三个：

  • 导出时记得设model.eval()，选合适的opset版本
  • 导出后一定要做兼容性检查，包括结构检查和数值精度对比
  • 遇到算子不兼容，优先考虑升级opset版本，或者把复杂操作拆成基础算子

下一章，我们会深入ONNX Runtime，看看怎么用它做高性能推理。到时候我会分享一些我在项目中用到的优化技巧，保证让你眼前一亮。

嗯，今天就到这儿。有问题随时找我。

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