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內部數據結構的理解, 操作起來不難。