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来自荷兰,面向未来
推荐阅读
vscode 如何全局替换
这里有一个向右的箭头点一下就从 “find” 变成 “replace” 了
ponponon阅读 86
又一款眼前一亮的Linux终端工具!
今天给大家介绍一款最近发现的功能十分强大,颜值非常高的一款终端工具。这个神器我是在其他公众号文章上看到的,但他们都没把它的强大之处介绍明白,所以我自己体验一波后,再向大家分享自己的体验。
良许赞 6阅读 1.9k
编译 PyTorch 模型
本篇文章译自英文文档 Compile PyTorch Models。作者是 Alex Wong。更多 TVM 中文文档可访问 →TVM 中文站。本文介绍了如何用 Relay 部署 PyTorch 模型。首先应安装 PyTorch。此外,还应安装 TorchVision,并将其...
超神经HyperAI赞 1阅读 93.1k
FastAPI性能碾压Flask?
不止一次的听过,FastAPI性能碾压Flask,直追Golang,不过一直没有测试过,今天闲着没事测试一下看看结果。不知道是哪里出了问题,结果大跌眼镜。
二毛erma0赞 2阅读 10.3k评论 3
Linux终端居然也可以做文件浏览器?
大家好,我是良许。在抖音上做直播已经整整 5 个月了,我很自豪我一路坚持到了现在【笑脸】最近我在做直播的时候,也开始学习鱼皮大佬,直播写代码。当然我不懂 Java 后端,因此就写写自己擅长的 Shell 脚本。但...
良许赞 1阅读 2.1k
Python之如何优雅的重试
为了避免偶尔的网络连接失败,需要加上重试机制,那么最简单的形式就是在对应的代码片段加一个循环,循环体里使用异常捕获,连接成功时退出循环,否则就重复执行相关逻辑,此时修改之后的函数f如下
Harpsichord1207赞 3阅读 7.4k
国产开源大模型盘点:ChatGLM-6B 和 MOSS 的应用与创新
ChatGLM-6B 是一个开源的、支持中英双语的对话语言模型,基于 General Language Model (GLM) 架构,具有 62 亿参数。结合模型量化技术,用户可以在消费级的显卡上进行本地部署(INT4 量化级别下最低只需 6GB 显存...
我的小熊不见了赞 1阅读 3.6k
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用
。你还可以使用@
来通知其他用户。