前言
构建onnx方式通常有两种:
1、通过代码转换成onnx结构,比如pytorch —> onnx
2、通过onnx 自定义结点,图,生成onnx结构
本文主要是简单学习和使用两种不同onnx结构,
下面以 Slice
结点进行分析
方式
方法一:pytorch --> onnx
暂缓,主要研究方式二
方法二: onnx
import onnx
from onnx import helper
from onnx import TensorProto
# 创建一个空的ONNX图
graph = helper.make_graph(nodes=[], name='Slice_Graph', inputs=[], outputs=[])
# 创建Slice节点的输入参数
input_data = helper.make_tensor_value_info('input_data', TensorProto.FLOAT, [3, 4, 5]) # 输入数据张量的形状为[3, 4, 5]
starts = helper.make_tensor('starts', TensorProto.INT64, [3], [0, 1, 2]) # 切片起始位置
ends = helper.make_tensor('ends', TensorProto.INT64, [3], [2, 3, 4]) # 切片结束位置
graph.input.extend([input_data])
graph.initializer.extend([starts, ends])
# 创建Slice节点
slice_node = helper.make_node('Slice', ['input_data', 'starts', 'ends'], ['output'], name='Slice_Node')
# 添加Slice节点到图中
graph.node.extend([slice_node])
# 创建Slice节点的输出参数
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [2, 2, 2]) # 输出张量形状为[2, 2, 2],表示切片后的子张量形状
graph.output.extend([output])
# 创建ONNX模型
model = helper.make_model(graph, producer_name='ONNX_Slice_Example')
# 保存ONNX模型到文件
onnx.save(model, 'slice_model.onnx')