ONNX初探

0x0. 背景

最近看了一些ONNX的资料，一个最大的感受就是这些资料太凌乱了。大多数都是在介绍ONNX模型转换中碰到的坑点以及解决办法。很少有文章可以系统的介绍ONNX的背景，分析ONNX格式，ONNX简化方法等。所以，综合了相当多资料之后我准备写一篇ONNX相关的文章，希望对大家有用。

0x1. 什么是ONNX？

简单描述一下官方介绍，开放神经网络交换（Open Neural Network Exchange）简称ONNX是微软和Facebook提出用来表示深度学习模型的开放格式。所谓开放就是ONNX定义了一组和环境，平台均无关的标准格式，来增强各种AI模型的可交互性。

换句话说，无论你使用何种训练框架训练模型（比如TensorFlow/Pytorch/OneFlow/Paddle），在训练完毕后你都可以将这些框架的模型统一转换ONNX这种统一的格式进行存储。注意ONNX文件不仅仅存储了神经网络模型的权重，同时也存储了模型的结构信息以及网络中每一层的输入输出和一些其它的辅助信息。我们直接从onnx的官方模型仓库拉一个yolov3-tiny的onnx模型（地址为：https://github.com/onnx/models/tree/master/vision/object_detection_segmentation/tiny-yolov3/model）用Netron可视化一下看看ONNX模型长什么样子。

这里我们可以看到ONNX的版本信息，这个ONNX模型是由Keras导出来的，以及模型的输入输出等信息，如果你对模型的输入输出有疑问可以直接看：https://github.com/onnx/models/blob/master/vision/object_detection_segmentation/tiny-yolov3/README.md。

在获得ONNX模型之后，模型部署人员自然就可以将这个模型部署到兼容ONNX的运行环境中去。这里一般还会设计到额外的模型转换工作，典型的比如在Android段利用NCNN部署模型，那么就需要将ONNX利用NCNN的转换工具转换到NCNN所支持的bin和param格式。

但在实际使用ONNX的过程中，大多数人对ONNX了解得并不多，仅仅认为它只是一个模型转换工具人而已，可以利用它完成模型转换和部署。正是因为对ONNX的不了解，在模型转换过程中出现的各种不兼容或者不支持让很多人浪费了大量时间。这篇文章将从理论和实践2个方面谈一谈ONNX。

0x2. ProtoBuf简介

在分析ONNX组织格式前我们需要了解ProtoBuf, 如果你非常了解ProtoBuf可以略过此节。
ONNX作为一个文件格式，我们自然需要一定的规则去读取我们想要的信息或者是写入我们需要保存信息。ONNX使用的是ProtoBuf这个序列化数据结构去存储神经网络的权重信息。熟悉Caffe或者Caffe2的同学应该知道，它们的模型存储数据结构协议也是Protobuf。这个从安装ONNX包的时候也可以看到：

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

Protobuf协议是一个以*.proto后缀文件为基础的，这个文件描述了用户自定义的数据结构。如果需要了解更多细节请参考后面的资料3，这里只是想表达ONNX是基于Protobuf来做数据存储和传输，那么自然onnx.proto就是ONNX格式文件了，接下来我们就分析一下ONNX格式。

0x3. ONNX格式分析

这一节我们来分析一下ONNX的组织格式，上面提到ONNX中最核心的部分就是onnx.proto（https://github.com/onnx/onnx/blob/master/onnx/onnx.proto）这个文件了，它定义了ONNX这个数据协议的规则和一些其它信息。现在是2020年1月，这个文件有700多行，我们没有必要把这个文件里面的每一行都贴出来，我们只要搞清楚里面的核心部分即可。在这个文件里面以message关键字开头的对象是我们需要关心的。我们列一下最核心的几个对象并解释一下它们之间的关系。

• ModelProto
• GraphProto
• NodeProto
• ValueInfoProto
• TensorProto
• AttributeProto

当我们加载了一个ONNX之后，我们获得的就是一个ModelProto，它包含了一些版本信息，生产者信息和一个GraphProto。在GraphProto里面又包含了四个repeated数组，它们分别是node(NodeProto类型)，input(ValueInfoProto类型)，output(ValueInfoProto类型)和initializer(TensorProto类型)，其中node中存放了模型中所有的计算节点，input存放了模型的输入节点，output存放了模型中所有的输出节点，initializer存放了模型的所有权重参数。

我们知道要完整的表达一个神经网络，不仅仅要知道网络的各个节点信息，还要知道它们的拓扑关系。这个拓扑关系在ONNX中是如何表示的呢？ONNX的每个计算节点都会有input和output两个数组，这两个数组是string类型，通过input和output的指向关系，我们就可以利用上述信息快速构建出一个深度学习模型的拓扑图。这里要注意一下，GraphProto中的input数组不仅包含我们一般理解中的图片输入的那个节点，还包含了模型中所有的权重。例如，Conv层里面的W权重实体是保存在initializer中的，那么相应的会有一个同名的输入在input中，其背后的逻辑应该是把权重也看成模型的输入，并通过initializer中的权重实体来对这个输入做初始化，即一个赋值的过程。

最后，每个计算节点中还包含了一个AttributeProto数组，用来描述该节点的属性，比如Conv节点或者说卷积层的属性包含group，pad，strides等等，每一个计算节点的属性，输入输出信息都详细记录在https://github.com/onnx/onnx/blob/master/docs/Operators.md。

0x4. onnx.helper

现在我们知道ONNX是把一个网络的每一层或者说一个算子当成节点node，使用这些Node去构建一个Graph，即一个网络。最后将Graph和其它的生产者信息，版本信息等合并在一起生成一个model，也即是最终的ONNX模型文件。
在构建ONNX模型的时候，https://github.com/onnx/onnx/blob/master/onnx/helper.py这个文件非常重要，我们可以利用它提供的make_node，make_graph，make_tensor等等接口完成一个ONNX模型的构建，一个示例如下：

import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto


# The protobuf definition can be found here:
# https://github.com/onnx/onnx/blob/master/onnx/onnx.proto


# Create one input (ValueInfoProto)
X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [3, 2])
pads = helper.make_tensor_value_info('pads', TensorProto.FLOAT, [1, 4])

value = helper.make_tensor_value_info('value', AttributeProto.FLOAT, [1])


# Create one output (ValueInfoProto)
Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [3, 4])

# Create a node (NodeProto) - This is based on Pad-11
node_def = helper.make_node(
    'Pad', # node name
    ['X', 'pads', 'value'], # inputs
    ['Y'], # outputs
    mode='constant', # attributes
)

# Create the graph (GraphProto)
graph_def = helper.make_graph(
    [node_def],
    'test-model',
    [X, pads, value],
    [Y],
)

# Create the model (ModelProto)
model_def = helper.make_model(graph_def, producer_name='onnx-example')

print('The model is:\n{}'.format(model_def))
onnx.checker.check_model(model_def)
print('The model is checked!')

这个官方示例为我们演示了如何使用onnx.helper的make_tensor，make_tensor_value_info，make_attribute，make_node，make_graph，make_node等方法来完整构建了一个ONNX模型。需要注意的是在上面的例子中，输入数据是一个一维Tensor，初始维度为[2]，这也是为什么经过维度为[1,4]的Pad操作之后获得的输出Tensor维度为[3,4]。另外由于Pad操作是没有带任何权重信息的，所以当你打印ONNX模型时，ModelProto的GraphProto是没有initializer这个属性的。

0x5. onnx-simplifier

原本这里是要总结一些使用ONNX进行模型部署经常碰到一些因为版本兼容性，或者各种框架OP没有对齐等原因导致的各种BUG。但是这样的话显得篇幅会很长，所以这里以一个经典的Pytorch转ONNX的reshape问题为例子，来尝试讲解一下大老师的onnx-simplifier，个人认为这个问题是基于ONNX模型部署最经典的问题，希望解决这个问题的过程中大家能有所收获。

问题发生在当我们想把下面这段代码导出ONNX模型时：