1. ONNX 简介

1.1 什么是 ONNX

开放神经网络交换 ONNX（Open Neural Network Exchange）是一套表示深度神经网络模型的开放格式，由微软和 Facebook 于 2017 推出，然后迅速得到了各大厂商和框架的支持。通过短短几年的发展，已经成为表示深度学习模型的实际标准，并且通过 ONNX-ML，可以支持传统非神经网络机器学习模型，大有一统整个 AI 模型交换标准。

1.2 ONNX 的核心思想

ONNX 定义了一组与环境和平台无关的标准格式，为 AI 模型的互操作性提供了基础，使 AI 模型可以在不同框架和环境下交互使用。硬件和软件厂商可以基于 ONNX 标准优化模型性能，让所有兼容 ONNX 标准的框架受益。目前，ONNX 主要关注在模型预测方面（inferring），使用不同框架训练的模型，转化为 ONNX 格式后，可以很容易的部署在兼容 ONNX 的运行环境中。

1.3 ONNX 的存储方式 —— ProtoBuf

ONNX 使用的是 Protobuf 这个序列化数据结构去存储神经网络的权重信息。

Protobuf 是一种轻便高效的结构化数据存储格式，可以用于结构化数据串行化，或者说序列化。它很适合做数据存储或数据交换格式。可用于通讯协议、数据存储等领域的语言无关、平台无关、可扩展的序列化结构数据格式。目前提供了 C++、Java、Python 三种语言的 API。

1.4 ONNX 组成部分 {Opset}

ONNX 规范由以下几个部分组成：

1. 一个可扩展的计算图模型：定义了通用的计算图中间表示法（Intermediate Representation）。
2. opset：ai.onnx 和 ai.onnx.ml。
   • ai.onnx 是默认的操作符集，主要针对神经网络模型
   • ai.onnx.ml 主要适用于传统非神经网络机器学习模型
3. 标准数据类型：包括张量（tensors）、序列（sequences）和映射（maps）。

opset：operator set，可以翻译为算子集合。

目前，ONNX 规范有两个官方变体，主要区别在与支持的类型和默认的操作符集（opset）。ONNX 神经网络变体只使用张量作为输入和输出；而作为支持传统机器学习模型的 ONNX-ML，还可以识别序列和映射，ONNX-ML 为支持非神经网络算法扩展了 ONNX 操作符集。

1.5 ONNX 主要协议

• ModelProto（模型协议）： 定义整个神经网络模型的结构，包括模型的元数据、图结构以及其他相关信息。
• GraphProto（图协议）： 描述神经网络的计算图结构，包括节点（NodeProto）、边（连接节点的边）等信息。
• NodeProto（节点协议）： 用于定义计算图中的节点，每个节点表示一个操作或计算步骤，包括该节点的输入、输出、操作类型等信息。
• ValueInfoProto（值信息协议）： 用于描述计算图中的值（如张量）的信息，包括名称、数据类型、形状等。
• TensorProto（张量协议）： 用于描述神经网络中的张量，包括张量的数据、形状、数据类型等信息。
• AttributeProto（属性协议）： 用于表示节点或图的属性，这些属性可能包含操作的参数、超参数等信息。

1.6 ONNX 的粒度与运行速度的关系

主流的模型部署有两种路径，以 TensorRT 为例，一种是 PyTorch->ONNX->TensorRT，另一种是 PyTorch->Caffe->TensorRT，两种转换路径的对比如下：

上面的表列了 ONNX 和 Caffe 的几点区别，其中最重要的区别就是 op 的粒度。举个例子，如果对 Bert 的 Attention 层做转换，ONNX 会把它变成 MatMul, Scale, SoftMax 的组合，而 Caffe 可能会直接生成一个叫做 Multi-Head Attention 的层，同时告诉 CUDA 工程师：“你去给我写一个大 kernel“（很怀疑发展到最后会不会把 ResNet50 都变成一个层 😂）

因此如果某天一个研究员提了一个新的 SOTA 的 op，很可能它直接就可以被转换成 ONNX（如果这个 op 在 PyTorch 的实现全都是用 Aten 的库拼接的），但是对于 Caffe 的工程师，需要重新写一个 kernel。

ATen 是 PyTorch 内置的 C++ 张量计算库，PyTorch 算子在底层绝大多数计算都是用 ATen 实现的。

讯享网

细粒度 op 的好处就是非常灵活，坏处就是速度会比较慢。这几年有很多工作都是在做 op fushion（比如把卷积和它后面的 ReLU 合到一起算），也就是把小 op 拼成大 op。

TensorRT 是 NVIDIA 推出的部署框架，自然性能是首要考量的，因此 Layer 的粒度都很粗（粗粒度代表着有大 op，从而速度会快）。在这种情况下把 Caffe 转换过去有天然的优势。

除此之外粗粒度也可以解决分支的问题。TensorRT 眼里的神经网络就是一个单纯的 DAG（有向无环图）：给定固定 shape 的输入，执行相同的运算，得到固定 shape 的输出。

在 评估一个自定义的节点 中有相关的实验。通过实验我们可以知道，将多个算子合在一起称之为 fusion，这个 fusion 是可以快加模型速度的。

2. ONNX 示例

2.1 线性回归（Linear Regression）{example1}

线性回归是机器学习中最简单的模型，由以下表达式描述：

$$Y=XA+B$$

我们可以将其看作是三个变量 $Y=f(X,A,B)$ 分解成 y = Add(MatMul(X, A), B)。这是我们需要用 ONNX 运算符表示的内容。首先是使用 ONNX 运算符实现一个函数。ONNX 是强类型的，必须为函数的输入和输出定义形状和类型。也就是说，我们需要四个函数来构建图，其中包括 make 函数：

• make_tensor_value_info：根据其形状和类型声明变量（输入或输出）
• make_node：创建由操作（op 类型）、其输入和输出定义的节点
• make_graph：创建一个带有前两个函数创建的对象的 ONNX 图的函数
• make_model：最后一个函数，将图和附加元数据合并

在整个创建过程中，我们需要为图的每个节点的每个输入和输出赋予一个名称。图的输入和输出由 ONNX 对象定义，使用字符串引用中间结果。下面是示例代码。

import onnx from onnx import TensorProto from onnx.helper import (make_model, make_node, make_graph, make_tensor, make_tensor_value_info) from onnx.checker import check_model # -------------------------- inputs -------------------------- # 'X'是名称，TensorProto.FLOAT是类型，[None, None]是形状。 X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None]) A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None]) B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None]) # -------------------------- outputs(形状未定义) -------------------------- Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None]) # -------------------------- nodes -------------------------- # 它创建一个由运算符类型MatMul定义的节点，'X'、'A'是节点的输入，'XA'是输出。 node1 = make_node(op_type='MatMul', inputs=['X', 'A'], outputs=['XA']) node2 = make_node(op_type='Add', inputs=['XA', 'B'], outputs=['Y']) # -------------------------- graph -------------------------- # 从节点到图，图是由节点列表、输入列表、输出列表和名称构建的。 graph = make_graph(nodes=[node1, node2], # 节点 name='lr', # 名称 inputs=[X, A, B], # 输入节点 outputs=[Y]) # 输出节点 # -------------------------- model -------------------------- # ONNX图，这种情况下没有元数据。 onnx_model = make_model(graph=graph) # 让我们检查模型是否一致，这个函数在“Checker and Shape Inference”部分有描述。 check_model(model=onnx_model) # 如果测试失败，将引发异常 print(onnx_model) # 将这个模型保存到本地 onnx.save_model(onnx_model, 'ONNX/saves/linear_regression.onnx') 

模型打印结果：

讯享网ir_version: 9 opset_import { version: 20 } graph { node { input: "X" input: "A" output: "XA" op_type: "MatMul" } node { input: "XA" input: "B" output: "Y" op_type: "Add" } name: "lr" input { name: "X" type { tensor_type { elem_type: 1 shape { dim { } dim { } } } } } input { name: "A" type { tensor_type { elem_type: 1 shape { dim { } dim { } } } } } input { name: "B" type { tensor_type { elem_type: 1 shape { dim { } dim { } } } } } output { name: "Y" type { tensor_type { elem_type: 1 shape { dim { } } } } } } 

⚠️ check_model() 这个函数的目的是检查模型是否一直，它没有返回值，如果模型有问题，那么这个函数会自动抛出异常。

我们用 Netron 看一下这个模型：