4.26 构建onnx结构模型-Suqeeze

2023-12-30 14:08:34

前言

构建onnx方式通常有两种:
1、通过代码转换成onnx结构,比如pytorch —> onnx
2、通过onnx 自定义结点,图,生成onnx结构

本文主要是简单学习和使用两种不同onnx结构,
下面以 Suqeeze 结点进行分析

在这里插入图片描述

方式

方法一:pytorch --> onnx

暂缓,主要研究方式二

方法二: onnx

import onnx
from onnx import helper
from onnx import TensorProto

# 创建一个squeeze节点
def create_squeeze_node(input_name, output_name, axes):
    node = helper.make_node(
        'Squeeze',
        inputs=[input_name],
        outputs=[output_name],
        axes=axes
    )
    return node

# 创建一个ONNX图
def create_onnx_graph():
    input_name = 'input'
    output_name = 'output'
    axes = [1, 2]  # 压缩的轴

    # 创建输入张量
    input_tensor = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, [1, 1, 3, 3])

    # 创建输出张量
    output_tensor = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, [3, 3])

    # 创建squeeze节点
    squeeze_node = create_squeeze_node(input_name, output_name, axes)

    # 创建ONNX图
    graph = helper.make_graph(
        [squeeze_node],
        'squeeze_graph',
        [input_tensor],
        [output_tensor]
    )

    # 创建ONNX模型
    model = helper.make_model(graph)

    return model

# 保存ONNX模型
def save_onnx_model(model, file_path):
    onnx.save(model, file_path)

# 创建并保存ONNX模型
model = create_onnx_graph()
save_onnx_model(model, 'squeeze_model.onnx')

文章来源:https://blog.csdn.net/weixin_45063703/article/details/135304271
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。