onnx 总结

简介

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

ProtoBuf 简介

ONNX 使用的是 Protobuf 这个序列化数据结构去存储神经网络的权重信息。Protobuf 是一种轻便高效的结构化数据存储格式，可以用于结构化数据串行化，或者说序列化。它很适合做数据存储或数据交换格式。可用于通讯协议、数据存储等领域的语言无关、平台无关、可扩展的序列化结构数据格式。目前提供了 C++、Java、Python 三种语言的 API。
Protobuf 协议是一个以 *.proto 后缀文件为基础的，这个文件描述了用户自定义的数据结构。

ONNX 格式分析

ONNX中最核心的部分就是 onnx.proto，这个文件定义了 ONNX 数据协议的规则和一些其它信息，比较重要的有：

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

加载了一个ONNX之后，我们获得的就是一个 ModelProto，它包含了一些版本信息，生产者信息和一个GraphProto，在 GraphProto 里面又包含了四个 repeated 数组，它们分别是：

• node(NodeProto类型):存放了模型中所有的计算节点
• input(ValueInfoProto类型):存放了模型所有的输入节点
• output(ValueInfoProto类型):存放了模型中所有的输出节点
• initializer(TensorProto类型):存放了模型的所有权重参数

要完整的表达一个神经网络，不仅仅要知道网络的各个节点信息，还要知道它们的拓扑关系。ONNX 的每个计算节点都会有 input 和 output 两个数组，这两个数组是 string 类型，通过 input 和 output 的指向关系，我们就可以利用上述信息快速构建出一个深度学习模型的拓扑图。GraphProto 中的 input 数组不仅包含我们一般理解中的图片输入的那个节点，还包含了模型中所有的权重。每个计算节点中还包含了一个 AttributeProto 数组，用来描述该节点的属性，比如 Conv 节点或者说卷积层的属性包含 group，pad，strides 等等。

netron 查看模型结构

目前共有 3 种方式：

• 第一种是以软件的方式下载安装 netron，然后打开软件载入模型
• 第二种是将 netron 作为 python 库进行安装，在 python 代码调用 netron 库来载入模型进行可视化。 可以通过 pip install netron 进行安装
• 第三种是使用在线的工具直接查看

以 tinyyolov2-7.onnx 为例：

在 netron 中，如果想看到特征图的维度，需要在导出为 onnx 的时候，同时加上特征图维度信息。这个操作需要 onnx 库的帮助，可以通过 pip install onnx 进行安装。具体代码如下：

import torchvision.models as models
import torch

import onnx
import onnx.utils
import onnx.version_converter


# 定义数据+网络
data = torch.randn(2, 3, 256, 256)
net = models.resnet34()

# 导出
torch.onnx.export(
    net,
    data,
    'model.onnx',
    export_params=True,
    opset_version=8,
)

# 增加维度信息
model_file = 'model.onnx'
onnx_model = onnx.load(model_file)
onnx.save(onnx.shape_inference.infer_shapes(onnx_model), model_file)