ONNX和ONNX Runtime

1 ONNX

ONNX是表示模型的一种格式, 它提供了对模型的一种中间表示。

1.1 ONNX's Design Principles:

• 既支持深度神经网络, 也支持传统的机器学习
• 可解释
• 后向兼容(backward compatible)
• compact and cross-platform representation for serialization

1.2 onnx.proto

ONNX选择使用Google的ProtoBuf(Protocol Buffers: 一种开源跨平台的序列化数据结构的协议, 也是一种序列化对象框架（或者说编解码框架)来表示计算图。.onnx文件中保存的就是计算图

可以参考onnx.proto来理解.onnx的数据结构。其大致组成可以分为以下3个部分:

• 计算图

import onnx
model = onnx.load(onnx_model_path)
graph = model.graph

# inputs和outputs存储的是tensor_value_info
inputs = graph.input
outputs = graph.output
# initializer存储了具体的weights值
initializers = graph.initializer
# node对应operator
nodes = graph.node

• 数据类型
  • Tensor Type
  • None-Tensor Type(Sequence && Map)
• operator

node = onnx.helper.make_node(
    'Relu',
    inputs=['x'],
    outputs=['y'],
)

1.3 dynamic batchsize

1.3.1 将pytorch模型转成dynamic batchsize的onnx模型

import torch
import torchvision
from torch.autograd import Variable

input_name = ['input']
output_name = ['output']
input = Variable(torch.randn(1, 3, 224, 224)).cuda()
model = torchvision.models.resnet50(pretrained=True).cuda()
dynamic_onnx = True
if dynamic_onnx:
    dynamic_axes = {
        'input': {0: "batch_size"},
        'output': {0: "batch_size"}
    }
    torch.onnx.export(model, input, 'resnet50_dynamic.onnx', input_names=input_name, output_names=output_name, dynamic_axes=dynamic_axes, verbose=True)
else:
    torch.onnx.export(model, input, 'resnet50.onnx', input_names=input_name, output_names=output_name, verbose=True)

1.3.2 将不是dynamic batchsize的onnx模型修改成dynamic batchsize形式的onnx模型

import onnx

def cvt2dynamic(src_onnx_model_path, dst_onnx_model_path):
    '''
    将不是dynamic batchsize的onnx模型转换成dynamic batchsize的onnx模型
    '''
    model = onnx.load(src_onnx_model_path)
    graph = model.graph
    inputs = graph.input
    outputs = graph.output
    # 通过netron可是化工具找到输入和输出的key, 比如在resnet50中, 这个键值就是input和output
    for inp in inputs:
        if inp.name == 'input':
            dim_c = inp.type.tensor_type.shape.dim[1].dim_value
            dim_h = inp.type.tensor_type.shape.dim[2].dim_value
            dim_w = inp.type.tensor_type.shape.dim[3].dim_value
            new_shape = ['batch_size', dim_c, dim_h, dim_w]
            inputs.remove(inp)
            new_inp = onnx.helper.make_tensor_value_info(inp.name, onnx.TensorProto.FLOAT, new_shape)
            inputs.extend([new_inp])

    for oup in outputs:
        if oup.name == 'output':
            dim_n = oup.type.tensor_type.shape.dim[1].dim_value
            new_shape = ['batch_size', dim_n]
            outputs.remove(inp)
            new_oup = onnx.helper.make_tensor_value_info(oup.name, onnx.TensorProto.FLOAT, new_shape)
            outputs.extend([new_oup])

    onnx.save(model, dst_onnx_model_path)

参考资料:
1 Creating and Modifying ONNX Model Using ONNX Python API
2 onnx model zoo提供了很多预训练好的onnx模型

2 ONNX Runtime

ONNX Runtime可以对ONNX模型进行推理。如下图所示，ONNX Runtime提供了以下功能:

• 计算图优化
• 模型划分: 将计算图划分成不同的部分, 可以放到不同的execution provider进行运算
• Execution provider: 各种硬件加速器

参考资料:
1 ONNX and ONNX-Runtime
2 ONNX and ONNX-Runtime Slides

3 ONNX Graphsurgeon

onnx graphsurgeon是NVIDIA提供的对onnx模型进行操作的工具，可以在github仓库中了解相关的例子, 有了对onnx内部数据结构的理解, 操作起来不难。