首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >Pytorch轻松学-RetinaNet模型导出与C++部署

Pytorch轻松学-RetinaNet模型导出与C++部署

作者头像
OpenCV学堂
发布2023-01-04 17:03:14
9070
发布2023-01-04 17:03:14
举报

微信公众号:OpenCV学堂

RetinaNet模型导出

在Pytorch的torchvision框架中支持对象检测模型主要包括:

-SSD
-Faster-RCNN
-Mask-RCNN
-FCOS
-RetinaNet
-KeyPointRCNN

亲测以上模型除了SSD导出ONNX格式无法被ONNXRUNTIME、OpenVINO2022解析之外,其他模型导出均可正常加载部署并推理使用。SSD导出无法使用Pytorch官方已经说了,是因为torchvision的Bug,但是好像还一直没有解决。RetinaNet支持自定义模型训练,这个我已经集成到OpenMV工具软件中,准备好数据集支持一键零代码训练自定义数据,导出ONNX模型。单独导出RetinaNet模型为ONNX的脚本如下:

model = tv.models.detection.retinanet_resnet50_fpn(pretrained=True)
dummy_input = torch.randn(1, 3, 1333, 800)
model.eval()
model(dummy_input)
im = torch.zeros(1, 3, 1333, 800).to("cpu")
torch.onnx.export(model, im,
                    "retinanet_resnet50_fpn.onnx",
                    verbose=False,
                    opset_version=11,
                    training=torch.onnx.TrainingMode.EVAL,
                    do_constant_folding=True,
                    input_names=['input'],
                    output_names=['boxes', "scores", "labels"],
                    dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'}}
                  )

导出之后可以查看模型输入输出格式如下:

输入支持动态图像大小,输出分别是boxes直接输出位置、scores是每个box的置信度,值在0~1之间、labels是类别标签索引值。

推理测试

分别实现了ONNXRUNTIME与OpenVINO2022推理C++代码,代码其实跟C++版本的YOLOv5+OpenVINO2022的代码类似,唯一不同的时候需要设置一下动态输入跟每次推理时候的图像实际大小,就是要reshape一下。而ONNXRUNTIME代码相对比较简单,直接传入图像数据即可。分别封装成两个类,集成到QT中演示效果如下:

然后这两天把QT的这个推理引擎项目重新整理一下目录结构,把UI跟infer分别放在不同的子目录中,修改后的项目目录结构如下:

这样看上去会好点,比较一目了然!请忽略项目名称,真的不止于支持YOLOv5推理!

01

版本兼容性问题

通过Pytorch1.7.1导出RetinaNet的ONNX模型可以部署到OpenVINO2022上正确推理,但是当我升级Pytorch版本从1.71到1.10版本之后,同样脚本导出的ONNX格式模型,OpenVINO2022就无法识别了,会报如下错误:

RuntimeError: [ UNEXPECTED ] CPU plug-in doesn't support If operation with dynamic rank. Operation name: 3217

然后我重新切换到pytorch1.7版本之后导出又可以加载并正常推理了;但是ONNXRUNTIME一直可以。

本文参与 腾讯云自媒体分享计划,分享自微信公众号。
原始发表:2022-12-16,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 OpenCV学堂 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体分享计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
相关产品与服务
对象存储
对象存储(Cloud Object Storage,COS)是由腾讯云推出的无目录层次结构、无数据格式限制,可容纳海量数据且支持 HTTP/HTTPS 协议访问的分布式存储服务。腾讯云 COS 的存储桶空间无容量上限,无需分区管理,适用于 CDN 数据分发、数据万象处理或大数据计算与分析的数据湖等多种场景。
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档