构建onnx方式通常有两种:
1、通过代码转换成onnx结构,比如pytorch —> onnx
2、通过onnx 自定义结点,图,生成onnx结构
本文主要是简单学习和使用两种不同onnx结构,
下面以 Expand
结点进行分析
暂缓,主要研究方式二
import onnx
from onnx import helper
import numpy as np
# 定义输入数据
input_data = np.array([1, 2, 3])
# 创建expand节点
node = helper.make_node(
'Expand',
inputs=['input_data', 'shape'],
outputs=['output_data']
)
# 创建ONNX图
graph = helper.make_graph(
[node],
'expand_graph',
inputs=[
helper.make_tensor_value_info('input_data', onnx.TensorProto.FLOAT, [3]),
helper.make_tensor_value_info('shape', onnx.TensorProto.INT64, [3])
],
outputs=[
helper.make_tensor_value_info('output_data', onnx.TensorProto.FLOAT, [3, 3])
]
)
# 创建ONNX模型
model = helper.make_model(graph)
# 保存ONNX模型
onnx.save(model, 'expand_model.onnx')