ONNX简介 ONNX （Open Neural Network Exchange）是一种多框架共用的，开放协议的神经网络交换格式。ONNX使用Protobuf二进制格式来序列化模型。

        ONNX协议首先由微软和Meta提出，它定义了一组和环境、平台均无关的标准格式，用于增强各种AI模型的可交互性。没有ONNX模型之前，采用各个框架训练的模型只能通过特定的转换工具进行转换。有了ONNX协议后，无论使用何种框架训练模型，在训练完毕后都可以将这些框架的模型统一转换为ONNX这种统一的格式进行存储。ONNX文件不仅存储了神经网络模型的权重，也存储了模型的结构信息以及网络中每一层的输入输出等信息。

一，构建ONNX模型

        onnx模型包含如下几个protobuf，定义如下。在工程实践中，导出一个ONNX模型就是导出一个ModelProto。ModelProtp中包含GraphProto、版本号等信息。GraphProto又包含了NodeProto类型的node，ValueInfoProto类型的input和output，TensorProto类型的initializer。其中，node包含了模型中的所有OP，input和output包含了模型的输入和输出，initializer包含了所有的权重信息。 每个计算节点node中还包含了一个AttributeProto数组，用来描述该节点的属性，比如Conv节点或者卷积层的属性包含group，pad，strides等等。

 python解析如下：

        如下构建一个简单的onnx模型，其中helper。make_node的node name需要按照onnx包含的常见算子定义，否则onnx.checker.check_model不能通过。

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

X=helper.make_tensor_value_info('X',TensorProto.FLOAT,[1,3,32,32]) #n,ci,h,w
W = helper.make_tensor_value_info('W', TensorProto.FLOAT, [64,3,3,3,]) #co,ci,kh,kw
# bias = helper.make_tensor_value_info('bias', AttributeProto.FLOAT, [64])

# 创建输出 (ValueInfoProto)
Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1,64,16,16])

# 创建节点 (NodeProto)
node_def = helper.make_node(
    'Conv', # node name,Pad
    inputs=['X','W'], # inputs
    outputs=['Y'], # outputs
    kernel_shape=[3,3],
    pads=[0,0,0,0], # attributes
)

# 创建图 (GraphProto)
graph_def = helper.make_graph(
    [node_def],
    'test-model',
    [X,W],
    [Y],
)

# 创建模型 (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.save_model(model_def,'./models/simple_test.onnx')

        使用netron工具可视化刚才保存的onnx模型，如图，点击conv的问号，会在右边出来onnx内部对这个node的解释和示例。

二，基于pytorch导出ONNX模型

        torch.onnx.export()函数的参数可以具体去看看定义，基础的需要传给net, input的tensor。

另外，可以自定义算子，导出onnx模型，后面更新。

import torch,onnx,torchvision
from torch.autograd import Variable
from torchvision.models.resnet import resnet18

net =resnet18(pretrained=True)
x=Variable(torch.randn(100,3,32,32),requires_grad=True)
torch.onnx.export(net,x,'./models/resnet18.onnx')

'''
#对export传更多自定义参数示例
torch.onnx.export(model, # 搭建的网络
    x, # 输入张量
    'model.onnx', # 输出模型名称
    input_names=["input"], # 输入命名
    output_names=["output"], # 输出命名
    dynamic_axes={'input':{0:'batch'}, 'output':{0:'batch'}}  # 动态轴
)
'''

可视化后大概如下： 

三，给导出的ONNX模型添加中间算子的infer shapes

        如二中导出的onnx模型，只有最上层输入的shape,如果想知道每一层数据流的变化，需要借助onnx.shape_inference, 读取上一步的onnx模型，再重新保存一遍。

import onnx
from onnx import load_model,save_model
from onnx.shape_inference import infer_shapes    
onnx_model = load_model(modelPath)
onnx_model1 = infer_shapes(onnx_model)
save_model(onnx_model1, fPath + '_final.onnx')

参考：

 [ONNX从入门到放弃] 3. ONNX形状推理 - 知乎