pytorch faster_rcnn转为onnx格式

最新推荐文章于 2024-05-22 13:47:11 发布

yuanjiaqi_k

最新推荐文章于 2024-05-22 13:47:11 发布

阅读量628

点赞数

分类专栏： pytorch

本文链接：https://blog.csdn.net/yuanjiaqi_k/article/details/125890764

版权

PyTorch ONNX 模型转换 Faster R-CNN 预训练模型

关键词由CSDN通过智能技术生成

pytorch 专栏收录该内容

16 篇文章 2 订阅

订阅专栏

import torch
import torchvision.models as models

torch_model = torch.load("faster_rcnn.pth") # pytorch模型加载

model = models.resnet50()
model.fc = torch.nn.Linear(2048, 4)
model.load_state_dict(torch_model,False)

batch_size = 1  #批处理大小
input_shape = (3, 244, 384)   #输入数据,改成自己的输入shape

# #set the model to inference mode
model.eval()

x = torch.randn(batch_size, *input_shape)	# 生成张量
export_onnx_file = "test.onnx"			# 目的ONNX文件名
torch.onnx.export(model,
                    x,
                    export_onnx_file,
                    opset_version=10,
                    do_constant_folding=True,	# 是否执行常量折叠优化
                    input_names=["input"],	# 输入名
                    output_names=["output"],	# 输出名
                    dynamic_axes={"input":{0:"batch_size"},  # 批处理变量
                                    "output":{0:"batch_size"}})