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

微信公众号：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一直可以。