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
评论