前言
构建onnx方式通常有两种:
1、通过代码转换成onnx结构,比如pytorch —> onnx
2、通过onnx 自定义结点,图,生成onnx结构
本文主要是简单学习和使用两种不同onnx结构,
下面以 Unsuqeeze
结点进行分析
方式
方法一:pytorch --> onnx
暂缓,主要研究方式二
方法二: onnx
import onnx
from onnx import numpy_helper
from onnx import helper
from onnx import TensorProto
# 构建unsqueeze操作节点
unsqueeze_node = helper.make_node(
'Unsqueeze', # 节点类型
inputs=['input'], # 输入
outputs=['output'], # 输出
axes=[1] # 指定的维度
)
# 构建ONNX图
graph = helper.make_graph(
nodes=[unsqueeze_node], # 节点
name='unsqueeze_graph', # 图的名称
inputs=[helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 3, 4, 5])], # 输入信息
outputs=[helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 1, 3, 4, 5])], # 输出信息
)
# 构建ONNX模型
model = helper.make_model(graph, producer_name='onnx-unsqueeze-example')
# 保存ONNX模型
onnx.save(model, 'unsqueeze_model.onnx')