4.27 构建onnx结构模型-Unsuqeeze

前言

构建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')