ONNX

5 分钟阅读

概述

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

Framework / Tool Installation Tutorial
Caffe apple/coremltools and onnx/onnxmltools Example
Caffe2 part of caffe2 package Example
Chainer chainer/onnx-chainer Example
Cognitive Toolkit (CNTK) built-in Example
CoreML (Apple) onnx/onnxmltools Example
Keras onnx/keras-onnx Example
LibSVM onnx/onnxmltools Example
LightGBM onnx/onnxmltools Example
MATLAB Deep Learning Toolbox Example
ML.NET built-in Example
MXNet (Apache) part of mxnet package docs github Example
PyTorch part of pytorch package Example1, Example2, export for Windows ML, Extending support
SciKit-Learn onnx/sklearn-onnx Example
SINGA (Apache) - Github(experimental) built-in Example
TensorFlow onnx/tensorflow-onnx Examples

来自:onnx tutorials