pytorch resnet50 模型转成 onnx 的代码样例,batch_size 是动态可调节的

pytorch resnet50 模型转成 onnx 的代码样例,batch_size 是动态可调节的

以下是一个将 PyTorch 的 ResNet50 模型转换为 ONNX 格式的示例代码,并使用可变批处理大小:

import torch
import torchvision
import onnx
import onnxruntime

# 创建 PyTorch ResNet50 模型实例
model = torchvision.models.resnet50(pretrained=True)

# 定义一个 PyTorch 张量来模拟输入数据
batch_size = 4  # 定义批处理大小
input_shape = (batch_size, 3, 224, 224)
input_data = torch.randn(input_shape)

# 将模型转换为 ONNX 格式
output_path = "resnet50.onnx"
torch.onnx.export(model, input_data, output_path,
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}})

# 使用 ONNX 运行时加载模型
session = onnxruntime.InferenceSession(output_path)

# 定义一个 ONNX 张量来模拟输入数据
new_batch_size = 8  # 定义新的批处理大小
new_input_shape = (new_batch_size, 3, 224, 224)
new_input_data = torch.randn(new_input_shape)

# 在 ONNX 运行时中运行模型
outputs = session.run(["output"], {"input": new_input_data.numpy()})

注意,在将模型导出为 ONNX 格式时,需要指定 input_names 和 output_names 参数来指定输入和输出张量的名称,以便在 ONNX 运行时中使用。此外,我们还需要使用 dynamic_axes 参数来指定批处理大小的动态维度。最后,在 ONNX 运行时中使用 session.run() 方法来运行模型。


python后端实战经验分享
Python来自荷兰,面向未来

Python 后端工程师可加:企鹅群:537131912

2.6k 声望
588 粉丝
0 条评论
推荐阅读
vscode 如何全局替换
这里有一个向右的箭头点一下就从 “find” 变成 “replace” 了

ponponon阅读 86

又一款眼前一亮的Linux终端工具!
今天给大家介绍一款最近发现的功能十分强大,颜值非常高的一款终端工具。这个神器我是在其他公众号文章上看到的,但他们都没把它的强大之处介绍明白,所以我自己体验一波后,再向大家分享自己的体验。

良许6阅读 1.9k

编译 PyTorch 模型
本篇文章译自英文文档 Compile PyTorch Models。作者是 Alex Wong。更多 TVM 中文文档可访问 →TVM 中文站。本文介绍了如何用 Relay 部署 PyTorch 模型。首先应安装 PyTorch。此外,还应安装 TorchVision,并将其...

超神经HyperAI1阅读 93.1k

FastAPI性能碾压Flask?
不止一次的听过,FastAPI性能碾压Flask,直追Golang,不过一直没有测试过,今天闲着没事测试一下看看结果。不知道是哪里出了问题,结果大跌眼镜。

二毛erma02阅读 10.3k评论 3

封面图
Linux终端居然也可以做文件浏览器?
大家好,我是良许。在抖音上做直播已经整整 5 个月了,我很自豪我一路坚持到了现在【笑脸】最近我在做直播的时候,也开始学习鱼皮大佬,直播写代码。当然我不懂 Java 后端,因此就写写自己擅长的 Shell 脚本。但...

良许1阅读 2.1k

Python之如何优雅的重试
为了避免偶尔的网络连接失败,需要加上重试机制,那么最简单的形式就是在对应的代码片段加一个循环,循环体里使用异常捕获,连接成功时退出循环,否则就重复执行相关逻辑,此时修改之后的函数f如下

Harpsichord12073阅读 7.4k

国产开源大模型盘点:ChatGLM-6B 和 MOSS 的应用与创新
ChatGLM-6B 是一个开源的、支持中英双语的对话语言模型,基于 General Language Model (GLM) 架构,具有 62 亿参数。结合模型量化技术,用户可以在消费级的显卡上进行本地部署(INT4 量化级别下最低只需 6GB 显存...

我的小熊不见了1阅读 3.6k

Python 后端工程师可加:企鹅群:537131912

2.6k 声望
588 粉丝
宣传栏