ONNX

ONNX

概述

ONNX，Open Neural Network Exchange。由于神经网络架构很多，如caffe、tensorflow、pytorch、mxnet等等，模型结构各式各样，onnx旨在将模型结构统一起来。

官方代码：ONNX

算子操作：Operators

环境要求：pip install onnx onnxruntime onnx-simplifier netron

Netron查看结构

可以本地用netron直接打开onnx查看，也可以用命令行如下：
netron xxx.onnx --host 172.xx.xx.xx
然后用网页打开查看

构建模型

import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto
import numpy as np
 
# Create one input and output (ValueInfoProto)
shape = (1, 3, 5, 5)
input = helper.make_tensor_value_info('input', TensorProto.FLOAT, list(shape))
output = helper.make_tensor_value_info(
    'output', TensorProto.FLOAT, list(shape))
 
# create weight filter
w = np.random.randn(3, 3, 3, 3).astype(np.float32)
filter_node_def = onnx.helper.make_node(
    'Constant',
    inputs=[],
    outputs=['filter'],
    value=onnx.helper.make_tensor(
        name='const_tensor',
        data_type=onnx.TensorProto.FLOAT,
        dims=w.shape,
        vals=w.flatten(),
    ),
)
 
# create weight bias
b = np.random.randn(3).astype(np.float32)
bias_node_def = onnx.helper.make_node(
    'Constant',
    inputs=[],
    outputs=['bias'],
    value=onnx.helper.make_tensor(
        name='const_tensor',
        data_type=onnx.TensorProto.FLOAT,
        dims=b.shape,
        vals=b.flatten(),
    ),
)
 
# create conv node
conv_node_def = onnx.helper.make_node(
    "Conv",
    inputs=['input', 'filter', 'bias'],
    outputs=['output'],
    kernel_shape=[3, 3],
    pads=[1, 1, 1, 1],
    strides=[1, 1],
    dilations=[1, 1],
    group=1,
)
graph_def = helper.make_graph(
    [filter_node_def, bias_node_def, conv_node_def],
    'conv-model',
    [input],
    [output],
)
 
model_def = helper.make_model(graph_def, producer_name='onnx-example')
onnx.checker.check_model(model_def)
onnx.save(model_def, 'example.onnx')

生成的模型结构如下：

模型推理

参考链接：ONNX Runtime example

import onnxruntime
import numpy as np
 
x = np.random.randn(1, 3, 5, 5).astype(np.float32)
 
session = onnxruntime.InferenceSession("example.onnx")
input = {"input":x}
output = session.run(None, input)
 
print(output)

模型优化

from onnxsim import simplify
model = onnx.load("example.onnx")
model_simple,_ = simplify(model)
onnx.save(model_simple, "simple.onnx")

查看Params和Flops

安装 pip install onnx-opcounter

# params数量
onnx_opcounter xxxxx.onnx
# mac数量
onnx_opcounter --calculate-macs xxxx.onnx

其他模型转ONNX

https://harmonyhu.com/2021/03/27/onnx/

https://github.com/onnx/tutorials