4.7 构建onnx结构模型-Transpose

最新推荐文章于 2025-04-30 23:15:02 发布

nsq_ai

最新推荐文章于 2025-04-30 23:15:02 发布

阅读量1.2k

点赞数 12

分类专栏：模型推理文章标签： onnx onnxruntime

本文链接：https://blog.csdn.net/weixin_45063703/article/details/134667728

版权

模型推理专栏收录该内容

54 篇文章

订阅专栏

本文介绍了将PyTorch模型转换为ONNX结构的方法，重点讲解了使用`Transpose`节点的示例，包括如何创建onnx模型、检查和运行onnx计算图，以及使用CPUExecutionProvider执行模型。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言

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

本文主要是简单学习和使用两种不同onnx结构，
下面以transpose 结点进行分析
在这里插入图片描述

方式

方法一：pytorch --> onnx

暂缓，主要研究方式二

方法二： onnx

import onnx 
from onnx import TensorProto, helper, numpy_helper

def run():
    print("run start....\n")
    perm = [0,2,1]
    Transpose = helper.make_node(
            "Transpose",
            perm=perm,
            inputs=["input"],
            outputs=["output"],
            name="Transpose",
        )
    graph = helper.make_graph(
        nodes=[Transpose],
        name="test_graph",
        inputs=[
            helper.make_tensor_value_info(
                "input", TensorProto.FLOAT, (16,397,80)
            )
        ],  # use your input
        outputs=[
            helper.make_tensor_value_info(
                "output", TensorProto.FLOAT, (16,80,397)
            )
        ],  # use your output
        initializer=[
            helper.make_tensor("perm", TensorProto.INT64, [len(perm)], perm),
        ],
        )

    op = onnx.OperatorSetIdProto()
    op.version = 11
    model = helper.make_model(graph, opset_imports=[op])
    print("run done....\n")

    return model

if __name__ == "__main__":
    model = run()
    onnx.save(model, "./test_transpose.onnx")

运行onnx

import onnx
import onnxruntime
import numpy as np


# 检查onnx计算图
def check_onnx(mdoel):
    onnx.checker.check_model(model)
    # print(onnx.helper.printable_graph(model.graph))

def run(model):
    print(f'run start....\n')
    session = onnxruntime.InferenceSession(model,providers=['CPUExecutionProvider'])
    input_name1 = session.get_inputs()[0].name  
    input_data1= np.random.randn(16,397,80).astype(np.float32)
    print(f'input_data1 shape:{input_data1.shape}\n')

    output_name1 = session.get_outputs()[0].name

    pred_onx = session.run(
    [output_name1], {input_name1: input_data1})[0]

    print(f'pred_onx shape:{pred_onx.shape} \n')

    print(f'run end....\n')


if __name__ == '__main__':
    path = "./test_transpose.onnx"
    model = onnx.load("./test_transpose.onnx")
    check_onnx(model)
    run(path)