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

OpenCV学堂 2023-08-22 17:47

【TI资料】专为高压系统设计的新型MCU 如何增强电动汽车的实时控制能力？

点击上方蓝字关注我们

微信公众号：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推理！

版本兼容性问题

通过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一直可以。

学会用ONNXRUNTIME C++部署YOLOv5与YOLOv8对象检测，实例分割，姿态评估模型，TorchVision框架下支持的Faster-RCNN，RetinaNet对象检测、MaskRCNN实例分割、Deeplabv3 语义分割模型等主流深度学习模型导出ONNX与C++推理部署，轻松解决Torchvision框架下模型训练到部署落地难题。

整个视频课程通过案例代码实战驱动，手把手系统化教学，帮助大家掌握ONNXRUNTIME API2 C++开发的各种技巧，学会图像分类、对象检测、语义分割、实例分割、pytorch自定义模型部署等ONNXRUNTIME C++版本的模型推理与解析技巧，课程思维导图如下：