4.30 构建onnx结构模型-TopK-编程知识

4.30 构建onnx结构模型-TopK

前言

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

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

在这里插入图片描述

方式

方法一：pytorch --> onnx

import torch
import onnx
import onnx.utils
from onnx import helper, TensorProto# 创建一个简单的PyTorch模型
class TopKModel(torch.nn.Module):def __init__(self, k):super(TopKModel, self).__init__()self.k = kdef forward(self, x):values, indices = torch.topk(x, self.k)return values, indices# 导出PyTorch模型为ONNX格式
model = TopKModel(k=3)
x = torch.randn(1, 5)
torch.onnx.export(model, x, "topk_model_pytorch.onnx", input_names=["input"], output_names=["values", "indices"])

方法二： onnx

import torch
import onnx
import onnx.utils
from onnx import helper, TensorProto# 创建TopK ONNX node
topk_node = onnx.helper.make_node('TopK',inputs=['input'],outputs=['values', 'indices'],k=3
)# 创建ONNX图
graph = helper.make_graph([topk_node],name="topk_graph",inputs=[helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 5])],outputs=[helper.make_tensor_value_info('values', TensorProto.FLOAT, [1, 3]),helper.make_tensor_value_info('indices', TensorProto.INT64, [1, 3])]
)# 创建ONNX模型
model = helper.make_model(graph, producer_name='onnx-topk-example')# 保存ONNX模型
onnx.save(model, 'topk_model.onnx')