4.33 构建onnx结构模型-Expand

nsq_ai

已于 2023-12-31 10:38:31 修改

阅读量787

点赞数 8

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

于 2023-12-31 10:37:35 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_45063703/article/details/135313832

版权

模型推理专栏收录该内容

54 篇文章

订阅专栏

本文介绍了如何通过PyTorch转换为ONNX格式以及使用ONNX自定义节点（如Expand结点）来构建ONNX模型，包括定义输入数据、创建节点、图和模型的步骤。

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

前言

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

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

在这里插入图片描述

方式

方法一：pytorch --> onnx

暂缓，主要研究方式二

方法二： onnx

import onnx
from onnx import helper
import numpy as np

# 定义输入数据
input_data = np.array([1, 2, 3])

# 创建expand节点
node = helper.make_node(
    'Expand',
    inputs=['input_data', 'shape'],
    outputs=['output_data']
)

# 创建ONNX图
graph = helper.make_graph(
    [node],
    'expand_graph',
    inputs=[
        helper.make_tensor_value_info('input_data', onnx.TensorProto.FLOAT, [3]),
        helper.make_tensor_value_info('shape', onnx.TensorProto.INT64, [3])
    ],
    outputs=[
        helper.make_tensor_value_info('output_data', onnx.TensorProto.FLOAT, [3, 3])
    ]
)

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

# 保存ONNX模型
onnx.save(model, 'expand_model.onnx')