3. ONNX标准详解:算子集、IR与版本演进

聊到模型转换,ONNX 是个绕不开的坎。我最早接触 ONNX 是在 2019 年,当时要把 PyTorch 的模型部署到移动端,折腾了好几个晚上。说实话,那时候 ONNX 的生态远没现在成熟,踩坑踩到怀疑人生。但到今天,它已经是模型交换的事实标准了。

这一节,咱们把 ONNX 拆开揉碎了讲。我会从它的核心设计——算子集和中间表示(IR)入手,再聊聊版本演进,最后给个客观的优缺点分析。嗯,都是实战中摸爬滚打出来的经验。

3.1 ONNX 的算子集与 IR

ONNX 的核心,说白了就两样东西:算子集(Opset)中间表示(IR)。前者定义了你能用哪些操作,后者定义了模型的结构怎么描述。

3.1.1 算子集(Opset)

算子集就是 ONNX 官方维护的一套标准操作。比如卷积、池化、全连接、激活函数这些。每个算子都有明确的输入输出定义、属性参数和计算逻辑。

我个人习惯把算子集想象成「乐高积木」。你手里的积木种类越多,能搭出的模型就越丰富。ONNX 的算子集目前覆盖了主流的深度学习操作,大概 200 多个算子。

关键点: 算子集是版本化的。每个版本会新增、修改或废弃一些算子。比如 opset 11 引入了 ScatterND,opset 13 改进了 Softmax 的行为。

举个例子,一个简单的 Conv 算子在 ONNX 中长这样:

// 伪代码表示
Conv(
  inputs: [X, W, B],   // 输入特征图、权重、偏置
  attributes: {
    dilations: [1, 1],
    group: 1,
    kernel_shape: [3, 3],
    pads: [1, 1, 1, 1],
    strides: [1, 1]
  },
  output: Y
)

我在项目中遇到过一个问题:PyTorch 的 F.interpolate 默认使用 align_corners=False,但 ONNX 的 Resize 算子在低版本 opset 里不支持这个参数。结果模型导出后,推理结果完全不对。后来我强制指定了 opset 11 以上才解决。

避坑指南: 导出 ONNX 时,尽量使用较新的 opset 版本(比如 opset 17 或 18)。但也要注意推理框架的支持情况。比如 TensorRT 8.x 最高只支持到 opset 13。

3.1.2 中间表示(IR)

IR 定义了模型的结构。它用 计算图(Graph) 来描述整个网络。一个 ONNX 模型包含:

  • 节点(Node):每个算子就是一个节点
  • 张量(Tensor):节点之间的数据流动
  • 初始器(Initializer):权重、偏置等常量数据
  • 输入/输出(Input/Output):模型的入口和出口

你想想看,这其实和 TensorFlow 的 GraphDef、PyTorch 的 TorchScript 是同一个思路。但 ONNX 的 IR 更简洁,也更通用。

下面这张图展示了 ONNX 模型的结构:

ONNX 模型结构图 输入 (Input) Conv 节点 ReLU 节点 Pool 节点 输出 (Output) 初始器 (权重) 输入/输出 算子节点 初始器

这张图展示了一个简单的 ONNX 计算图。输入经过 Conv、ReLU、Pool 三个节点,最终得到输出。权重数据存储在初始器中,通过虚线连接到 Conv 节点。

3.2 ONNX 的版本演进

ONNX 从 2017 年开源到现在,经历了十几个大版本。我挑几个关键节点说说:

版本 发布时间 关键变化
opset 1-7 2017-2018 基础算子集,支持 CNN、RNN 等常见结构
opset 9 2019 新增 ConstantOfShapeEyeLike 等,完善动态形状支持
opset 11 2020 引入 ScatterNDGatherND,强化稀疏操作
opset 13 2021 Softmax 行为变更,新增 IfLoop 控制流算子
opset 15 2022 新增 BernoulliTrilu,优化量化支持
opset 18 2023 新增 GroupNormalizationLayerNormalization,完善 Transformer 支持

为什么会这样演进?说白了,深度学习技术在变,ONNX 得跟上。比如 Transformer 火了之后,ONNX 在 opset 18 里专门加了 LayerNorm 和 GroupNorm 的原生支持。以前这些都得用组合算子拼,性能差不说,还容易出错。

注意: 版本兼容性是个大坑。我曾经把一个 opset 13 的模型拿到只支持 opset 11 的推理引擎上跑,结果直接报错。建议导出时先确认目标平台的 opset 上限。

3.3 ONNX 的优缺点分析

聊完技术细节,咱们客观说说 ONNX 的好与不好。毕竟没有银弹,每个工具都有它的脾气。

优点

  • 生态广泛:几乎所有主流框架都支持导出 ONNX。PyTorch、TensorFlow、PaddlePaddle、MXNet... 你想到的它基本都支持。
  • 标准化程度高:算子定义清晰,IR 结构统一。不同框架之间的模型交换变得可行。
  • 工具链完善:ONNX Runtime 可以直接推理,ONNX-TensorRT 可以转成 TensorRT 引擎,ONNX-MLIR 可以编译到硬件后端。
  • 社区活跃:微软、Facebook、Intel、AMD 等大厂都在贡献。问题修复和新特性迭代很快。

缺点

  • 算子覆盖不全:虽然已经有 200+ 算子,但总有些骚操作 ONNX 不支持。比如 PyTorch 的 torch.einsum,ONNX 就没有原生算子,得用组合方式替代。
  • 版本兼容性差:不同 opset 之间不保证向前兼容。你导出的模型换个环境可能就跑不起来了。
  • 性能不是最优:ONNX 本身是个中间表示,不是为推理优化的。直接跑 ONNX Runtime 的性能,通常不如专门优化的推理引擎(比如 TensorRT、OpenVINO)。
  • 调试困难:模型转换出错时,错误信息往往很模糊。我经常得一行一行对比原始模型和 ONNX 模型的输出,才能定位问题。
我的建议: 把 ONNX 当作「中间桥梁」来用,而不是最终部署格式。先用 ONNX 完成模型交换,再转成目标平台的专用格式。这样既利用了 ONNX 的生态优势,又避免了它的性能短板。

嗯,关于 ONNX 的核心内容就这些。记住它的算子集和 IR 是基础,版本演进是你要关注的坑,优缺点分析能帮你做技术选型。下一节咱们聊聊更具体的转换工具和实战技巧。


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