TensorRT中Detectron2 Mask RCNN模型转换与优化实践

背景介绍

在计算机视觉领域，Mask RCNN是一种广泛使用的实例分割模型。本文将详细介绍如何将Facebook Research的Detectron2框架中的Mask RCNN模型转换为TensorRT格式，并解决转换过程中遇到的各种技术挑战。

模型转换流程

初始转换步骤

首先需要将Detectron2的Mask RCNN模型(pkl格式)转换为ONNX格式。这一步骤使用Detectron2提供的export_model.py脚本完成。值得注意的是，原始转换后的ONNX模型在ONNX Runtime环境下运行良好，能够正确检测目标。

ONNX到TensorRT的转换

接下来使用TensorRT提供的create_onnx.py脚本对ONNX模型进行进一步处理。这一步骤会引入TensorRT特有的插件(plugins)，导致模型无法再被ONNX Runtime直接运行。这是因为：

1. 转换后的ONNX IR版本升级到10，而ONNX Runtime最高只支持到版本9
2. TensorRT插件是专为TensorRT优化的自定义操作，ONNX Runtime无法识别

使用TensorRT 8.4.2-1版本成功将模型转换为TRT引擎，在NVIDIA T600 GPU上实现了显著的性能提升：

• 原始Detectron2模型：4 FPS
• TensorRT FP32模式：7 FPS
• TensorRT FP16模式：平均10 FPS

关键问题与解决方案

检测结果丢失问题

初始转换后模型几乎无法检测到任何目标。经过排查发现，问题出在模型输入尺寸设置上。正确的做法是：

1. 使用正确的配置文件：detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml
2. 在导出模型时添加1344x1344输入尺寸的特殊修改

TensorRT插件的重要性

TensorRT插件在模型转换中扮演两个关键角色：

1. 性能优化：自定义内核显著提升推理速度
2. 模型可运行性：原始ONNX图过于复杂，包含大量动态排名的条件分支，TensorRT编译器难以处理。插件可以简化这些结构

扩展到Keypoint RCNN模型

基于Mask RCNN的转换经验，成功将Keypoint RCNN模型也移植到TensorRT。关键点包括：

1. 修改ROI heads处理逻辑，正确处理关键点热图
2. 在后期处理中使用numpy实现热图到关键点的转换
3. 验证热图与边界框的对应关系

最佳实践建议

1. 版本控制：使用较旧的Detectron2版本(1年以上)可获得更好的兼容性
2. 工具链选择：优先使用TensorRT Python API而非trtexec，便于调试
3. 性能调优：根据实际需求调整NMS阈值，平衡精度与速度
4. 自定义开发：学习编写CUDA内核和TensorRT插件可大幅提升性能

环境配置建议

推荐使用NVIDIA官方提供的Docker镜像作为基础环境：

• 基础镜像：nvcr.io/nvidia/pytorch:22.08-py3
• ONNX GraphSurgeon版本：0.5.2(避免使用Git版本)

总结

通过系统的模型转换流程和问题排查，成功将Detectron2的Mask RCNN和Keypoint RCNN模型部署到TensorRT环境，实现了显著的性能提升。这一过程不仅涉及格式转换，更需要深入理解模型结构和TensorRT优化原理。希望本文的经验能为类似项目的实施提供有价值的参考。