前言
构建onnx方式通常有两种:
1、通过代码转换成onnx结构,比如pytorch —> onnx
2、通过onnx 自定义结点,图,生成onnx结构
本文主要是简单学习和使用两种不同onnx结构,
下面以 Pad
结点进行分析
方式
方法一:pytorch --> onnx
暂缓,主要研究方式二
方法二: onnx
import onnx
from onnx import helper
from onnx import TensorProto
# 创建一个空的ONNX图
graph = helper.make_graph(nodes=[], name='Pad_Graph', inputs=[], outputs=[])
# 创建一个输入节点
input = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 3, 32, 32])
graph.input.extend([input])
# 创建Pad节点
pad_value = 0.0 # 填充值
pads = [0, 0, 1, 1, 2, 2, 0, 0] # 填充的维度和数量
mode = b'constant' # 填充模式
pad_node = helper.make_node('Pad', ['input'], ['output'], name='Pad_Node', pads=pads, value=pad_value, mode=mode)
# 添加Pad节点到图中
graph.node.extend([pad_node])
# 创建输出节点
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 3, 36, 36])
graph.output.extend([output])
# 创建ONNX模型
model = helper.make_model(graph, producer_name='ONNX_Pad_Example')
# 保存ONNX模型到文件
onnx.save(model, 'pad_model.onnx')