FasterCNNN+ROIAlign+FPN在TensorRT上的部署的解决方案
Repo(欢迎star): https://github.com/thb1314/tensorrt-onnx-fasterrcnn-fpn-roiailigin FasterRCNN作为经典的two-staged detector，至今为止其精度配合FPN+ROIAlign依然SOTA。 然而FasterRCNN在TensorRT的部署却没有one-stage的检测器那么容易。因为TensorRT输入输出要求shape是确定值，就算新的版本支持动态batch_size，但是在一开始也要确定最大输入的batch_size，中间的算子是不允许运行时更改维度的。 这成为FasterRCNN+FPN+ROIAlign在TensorRT部署的难点，因为rpn的输出需要经过nms+roialign，尽管nms和roialign可以通过plugin的方式来实现，但是nms的输出是不固定的。 此外，在计算roialign计算时，应该采用fpn哪个输出层的feature是根据rpn输出的bbox来计算的，onnx对于这里的计算步骤用计算图的表示引入了大量的判断算子，这在TensorRT的onnx解析器中是不支持的。 本项目采用
torchvision.models.detection.fasterrcnn_resnet50_fpn为待部署模型，着手解决部署中产生的一些问题。 主要思路：将rpn部分和检测头分开，分成两个模型。rpn部分采用onnx解析，nms+roialigin部分使用原生cuda来做。处理后的feature送入header部分，然后再一次经过cuda实现的nms得到最终结果。
ONNX生成步骤
代码见repo下的model和test
- 将FasterRCNN中的rpn和header部分剥离开，从torchvision中移植
- 分别生成rpn和header部分的onnx文件
- 整理rpn和header的输入与输出，rpn的输入为图像，输出为topk后的bbox和fpn各层的输出。header的输入为roialign后的feature和bbox，输出为nms前的最终的bbox和score
上面步骤涉及到计算图的修建，过程较为繁琐。
TensorRT实现部分
tensorRT部分需要完成哪些事？
- 对输入进行预处理
- 对rpn的输出进行nms处理，以及求取roi_align后的feature，这里涉及到对每一个bbox应该在哪一层fpn输出feature上进行计算roialign
- 创造header的输入，将nms后的bbox和roi_align计算后的feature送入header部分
- 对header的输出进行nms处理
roi_align如何实现？
- 移植onnxruntime的cuda实现代码（验证导出onnx的时候，已经证明onnxruntime实现的正确性）
- 看原版论文和torchvision的实现计算fpn_level
- 将以上两点结合，对不同fpn_level动态选择不同feature进行roi_align计算
以上需求的实现可以在tensorrt_code/application/src/application/app_fasterrcnn中可以找到
仓库代码持续更新中，如果有问题可以在本文下面留言，也可以在仓库issue区提出问题，如果可以动动小手给个star，我将不胜感激。
创建日期: September 17, 2024