4.30 构建onnx结构模型-TopK

本文介绍了将PyTorch模型转换为ONNX格式以及通过ONNX自定义结点创建TopK模型的过程,展示了两种不同的方式来构建ONNX模型。

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

前言

构建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 = k

    def 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')

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

nsq_ai

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值