4.34 构建onnx结构模型-Gemm

前言

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

本文主要是简单学习和使用两种不同onnx结构,
下面以 Gemm 结点进行分析
4.34 构建onnx结构模型-Gemm_第1张图片

方式

方法一:pytorch --> onnx

暂缓,主要研究方式二

方法二: onnx

import onnx
from onnx import helper
from onnx import TensorProto

# 创建一个空的ONNX图
graph = helper.make_graph(nodes=[], name='Gemm_Graph', inputs=[], outputs=[])

# 创建Gemm节点的输入参数
input_data_A = helper.make_tensor_value_info('input_data_A', TensorProto.FLOAT, [3, 4])  # 输入数据张量A的形状为[3, 4]
input_data_B = helper.make_tensor_value_info('input_data_B', TensorProto.FLOAT, [4, 5])  # 输入数据张量B的形状为[4, 5]
input_data_C = helper.make_tensor_value_info('input_data_C', TensorProto.FLOAT, [3, 5])  # 输入数据张量C的形状为[3, 5]

alpha = 1.0  # Gemm操作中的alpha参数
beta = 1.0  # Gemm操作中的beta参数
transA = 0  # 是否对输入A进行转置
transB = 1  # 是否对输入B进行转置

graph.input.extend([input_data_A, input_data_B, input_data_C])

# 创建Gemm节点
gemm_node = helper.make_node('Gemm', ['input_data_A', 'input_data_B', 'input_data_C'], ['output'], alpha=alpha, beta=beta, transA=transA, transB=transB, name='Gemm_Node')  # 创建一个Gemm节点

# 添加Gemm节点到图中
graph.node.extend([gemm_node])

# 创建Gemm节点的输出参数
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [3, 5])  # 输出张量形状为[3, 5]
graph.output.extend([output])

# 创建ONNX模型
model = helper.make_model(graph, producer_name='ONNX_Gemm_Example')

# 保存ONNX模型到文件
onnx.save(model, 'gemm_model.onnx')

你可能感兴趣的:(模型推理,onnx,python)