工业检测模型ONNX转换兼容与精度调试实战 - 第1节

一、ONNX 到底是什么?

先说说我的理解。ONNX 的全称是 Open Neural Network Exchange,直译过来就是「开放神经网络交换格式」。但说白了,它就是一个中间格式,一个模型界的「通用语言」。

你想想看,平时我们用 PyTorch 训练模型,用 TensorFlow 训练模型,用 Keras、MXNet、PaddlePaddle……每个框架都有自己的模型保存格式。PyTorch 是 .pt 或 .pth,TensorFlow 是 .pb 或 .h5,Keras 是 .h5。这些格式互相不认,就像你拿着一张人民币去英国买面包,人家不收。

ONNX 就是那个「货币兑换处」。你把 PyTorch 模型转成 ONNX,再把 ONNX 转成 TensorFlow,或者直接拿 ONNX 去推理。一次转换,到处运行。

核心定义:ONNX 是一种开放、跨平台、跨框架的深度学习模型表示标准。它定义了一套统一的算子集和计算图格式,让不同框架之间可以自由交换模型。

1.1 ONNX 的诞生背景

2017 年,微软和 Facebook(现 Meta)联合推出了 ONNX。当时我还在做算法落地,印象特别深。那会儿团队里有人用 PyTorch 做研究,有人用 TensorFlow 做部署,每次模型交接都像在搞外交——先导出权重,再手动重写网络结构,再对齐精度,折腾好几天。

ONNX 的出现,就是为了解决这个痛点。它把模型的计算图、权重、参数、输入输出信息全部打包成一个 .onnx 文件。你只需要在 PyTorch 里调用 torch.onnx.export(),就能得到一个标准的 ONNX 模型。然后这个模型可以在 ONNX Runtime、TensorRT、OpenVINO、CoreML 等推理引擎上直接跑。

嗯,这里要注意:ONNX 本身不是推理框架,它只是一个模型表示格式。真正跑推理的是 ONNX Runtime 或者别的引擎。这个区别很多人一开始会搞混。

二、为什么需要 ONNX?

我总结了三层理由,从浅到深。

2.1 打破框架壁垒

这是最直接的价值。没有 ONNX 的时候,模型从研究到部署,通常要走这么一条路:

  1. 研究员用 PyTorch 训练模型
  2. 工程师把 PyTorch 模型转成 TensorFlow 或 Caffe
  3. 再转成 TensorRT 或 OpenVINO 能吃的格式
  4. 每一步都可能丢精度、改结构、调参数

有了 ONNX,流程变成:

  1. PyTorch 训练 → 导出 ONNX
  2. ONNX 直接给 ONNX Runtime / TensorRT / OpenVINO 用

一步到位。我在项目中遇到过最极端的情况:一个用 PyTorch 训练的语义分割模型,要部署到华为昇腾芯片上。没有 ONNX 的话,得先转成 Caffe,再转成昇腾的离线模型,中间精度掉了 3 个点。后来直接用 ONNX 转,精度只掉了 0.2 个点。

2.2 硬件加速的「通用入口」

现在主流的推理硬件——NVIDIA GPU、Intel CPU、ARM、NPU、FPGA——几乎都支持 ONNX。你只要把模型转成 ONNX,就能在绝大多数硬件上跑推理。

为什么会这样?因为硬件厂商不想为每个框架都写一套适配器。他们只需要支持 ONNX 这一种格式,就能覆盖 PyTorch、TensorFlow、PaddlePaddle 等所有主流框架的用户。这是典型的「多对一」模式。

推理引擎 支持的框架格式 ONNX 支持情况
ONNX Runtime ONNX 原生 ✅ 原生支持
TensorRT ONNX、TensorFlow、Caffe ✅ 推荐入口
OpenVINO ONNX、TensorFlow、MXNet ✅ 推荐入口
CoreML ONNX、TensorFlow、PyTorch ✅ 支持
NCNN ONNX、Caffe ✅ 支持
TFLite TensorFlow、ONNX ✅ 通过转换支持

2.3 模型部署的「标准化」

我个人习惯把 ONNX 看作模型部署的「接口契约」。你训练完模型,导出 ONNX,就等于定义了一个稳定的接口。后续的量化、剪枝、蒸馏、优化,都可以在 ONNX 层面做,而不需要动原始训练代码。

我曾经维护过一个项目,模型迭代了 5 个版本,每次都要重新写部署代码。后来我们定了个规矩:所有模型必须导出 ONNX,部署只认 ONNX。从那以后,部署代码几乎没改过,只换 .onnx 文件就行。

个人建议:如果你的项目涉及模型部署,尽早把 ONNX 纳入流程。哪怕你暂时只用 PyTorch,也建议养成导出 ONNX 的习惯。万一哪天要换硬件或框架,你会感谢自己这个决定。

三、ONNX 生态全景图

ONNX 不是孤立的。它周围有一整套工具链和生态。我画了一张图,帮你快速建立全局认知。

ONNX 生态全景图 ONNX 模型 训练框架 PyTorch TensorFlow Keras MXNet PaddlePaddle 推理引擎 ONNX Runtime TensorRT OpenVINO CoreML NCNN 工具链 ONNX 算子集 ONNX 转换工具 ONNX 优化工具 ONNX 量化工具 生态核心逻辑 训练框架 → ONNX 模型 → 推理引擎 | 工具链贯穿全流程

3.1 训练框架层

几乎所有主流深度学习框架都支持导出 ONNX。我列一下常用的:

  • PyTorch:通过 torch.onnx.export() 导出,支持动态/静态 batch
  • TensorFlow:通过 tf2onnx 工具转换
  • Keras:通过 keras2onnx 转换
  • MXNet:原生支持导出 ONNX
  • PaddlePaddle:通过 paddle2onnx 转换

这里有个坑:不同框架的算子实现细节不同,转换时可能会遇到不支持的算子。比如 PyTorch 的 torch.nn.functional.grid_sample,在 ONNX 的早期版本里就没有对应的算子。遇到这种情况,要么换实现方式,要么用 ONNX 的自定义算子扩展。

3.2 推理引擎层

ONNX 模型最终要在推理引擎上跑。我按使用场景分了三类:

场景 推荐引擎 特点
通用 CPU/GPU 推理 ONNX Runtime 官方引擎,支持最全,跨平台
NVIDIA GPU 加速 TensorRT 极致性能,支持 INT8 量化
Intel CPU/GPU 加速 OpenVINO Intel 硬件优化,支持边缘设备
Apple 设备 CoreML iPhone/iPad/Mac 原生加速
移动端/嵌入式 NCNN / MNN / TNN 轻量级,无依赖

我个人最常用的是 ONNX Runtime + TensorRT 的组合。先用 ONNX Runtime 做精度验证,再用 TensorRT 做性能优化。这个流程我在后面的章节会详细讲。

3.3 工具链层

ONNX 的工具链非常丰富,我挑几个最常用的:

  • onnx:核心库,提供模型加载、保存、检查、修改功能
  • onnx-simplifier:简化 ONNX 模型,去掉冗余节点
  • onnxoptimizer:优化计算图,合并算子、消除死节点
  • onnxruntime:推理引擎,也提供模型转换和量化工具
  • onnx2tf / tf2onnx:ONNX 和 TensorFlow 互转
  • polygraphy:NVIDIA 的工具,用于 ONNX 模型调试和精度对比

避坑指南:我曾经在转换一个 YOLOv5 模型时,发现 ONNX 模型里多了很多 Shape 和 Gather 节点,导致推理速度变慢。后来用 onnx-simplifier 一跑,节点数从 800 降到了 300,推理速度提升了 40%。所以,导出 ONNX 后一定要做简化,这是标准流程。

四、ONNX 的核心概念

最后,快速过一下 ONNX 的几个核心概念。这些在后面做转换和调试时会反复用到。

4.1 计算图(Graph)

ONNX 模型本质上是一个有向无环图(DAG)。图中的节点是算子(Operator),边是张量(Tensor)。输入从图的入口进入,经过一系列算子计算,最终从出口输出。

4.2 算子集(Opset)

ONNX 定义了一套标准算子,比如 Conv、Relu、MatMul、Softmax 等。每个版本会新增或修改算子。目前最新的 opset 是 21(2024 年)。

转换时要注意 opset 版本。比如 PyTorch 默认导出 opset 17,但 TensorRT 8.x 只支持到 opset 13。版本不匹配会导致转换失败。

4.3 元数据(Metadata)

ONNX 模型里还包含模型的元信息:模型名称、版本、作者、输入输出的名称和形状、数据类型等。这些信息在部署时非常有用。

小技巧:用 netron.app 打开 .onnx 文件,可以可视化查看计算图结构。我每次转换完模型,第一件事就是用 Netron 看一眼,确认结构对不对。这个习惯帮我发现过好几次转换错误。