MMDet 转 TensorRT

[!NOTE]

主分支用于支持 MMDetection>=3.0 的模型转换。 如果您希望在较旧版本的 MMDetection 上进行模型转换，请切换到以下分支：

• mmdet2trt=v0.5.0
• torch2trt_dynamic=v0.5.0
• amirstan_plugin=v0.5.0。

新闻

• 2024年2月：支持 MMDetection>=3.0

简介

该项目旨在支持使用 TensorRT 对 MMDetection 中的模型进行端到端部署。

掩码支持目前处于 实验 阶段。

功能：

• fp16
• int8（实验）
• 批量输入
• 动态输入形状
• 不同模块的组合
• DeepStream

要求

• 安装 MMDetection：

  pip install openmim
  mim install mmdet==3.3.0
  

• 安装 torch2trt_dynamic：

  git clone https://github.com/grimoire/torch2trt_dynamic.git torch2trt_dynamic
  cd torch2trt_dynamic
  pip install -e .
  

• 安装 amirstan_plugin：

  • 安装 TensorRT：TensorRT

  • 克隆仓库并构建插件

    git clone --depth=1 https://github.com/grimoire/amirstan_plugin.git
    cd amirstan_plugin
    git submodule update --init --progress --depth=1
    mkdir build
    cd build
    cmake -DTENSORRT_DIR=${TENSORRT_DIR} ..
    make -j10
    

    [!NOTE]

    请勿忘记 在 ~/.bashrc 中设置环境变量：

    export AMIRSTAN_LIBRARY_PATH=${amirstan_plugin_root}/build/lib
    

安装

主机

git clone https://github.com/grimoire/mmdetection-to-tensorrt.git
cd mmdetection-to-tensorrt
pip install -e .

Docker

构建 Docker 镜像

sudo docker build -t mmdet2trt_docker:v1.0 docker/

运行（将显示 CLI 入口点的帮助信息）

sudo docker run --gpus all -it --rm -v ${your_data_path}:${bind_path} mmdet2trt_docker:v1.0

或者，如果您想在容器内打开终端：

sudo docker run --gpus all -it --rm -v ${your_data_path}:${bind_path} --entrypoint bash mmdet2trt_docker:v1.0

示例转换：

sudo docker run --gpus all -it --rm -v ${your_data_path}:${bind_path} mmdet2trt_docker:v1.0 ${bind_path}/config.py ${bind_path}/checkpoint.pth ${bind_path}/output.trt

使用方法

从 mmdet 模型创建 TensorRT 模型。详细信息请参阅 getting_started.md。

CLI

# 转换可能需要几分钟。
mmdet2trt ${CONFIG_PATH} ${CHECKPOINT_PATH} ${OUTPUT_PATH}

运行 mmdet2trt -h 可获取可选参数的帮助信息。

Python

shape_ranges=dict(
    x=dict(
        min=[1,3,320,320],
        opt=[1,3,800,1344],
        max=[1,3,1344,1344],
    )
)
trt_model = mmdet2trt(cfg_path,
                      weight_path,
                      shape_ranges=shape_ranges,
                      fp16_mode=True)

# 保存转换后的模型
torch.save(trt_model.state_dict(), save_model_path)

# 如果您想在 C++ API 中使用该引擎，可以将其保存下来
with open(save_engine_path, mode='wb') as f:
    f.write(trt_model.state_dict()['engine'])

[!NOTE]

引擎的输入是 预处理后 的张量。 引擎的输出为 num_dets, bboxes, scores, class_ids。如果您启用了 enable_mask 标志，则还会有一个额外的输出 mask。 引擎输出的边界框未除以 scale_factor。

如何使用转换后的模型进行推理：

from mmdet.apis import inference_detector
from mmdet2trt.apis import create_wrap_detector

# 创建包装检测器
trt_detector = create_wrap_detector(trt_model, cfg_path, device_id)

# 结果格式与 MMDetection 相同
result = inference_detector(trt_detector, image_path)

您可以尝试 demo/inference.py 中的演示，或如果想使用 C++ API 进行推理，可以查看 demo/cpp。

更多详细信息请阅读 getting_started.md。

工作原理

大多数其他项目采用 PyTorch => ONNX => TensorRT 的流程，而本仓库则直接将 PyTorch 转换为 TensorRT，避免了不必要的 ONNX IR。有关详细信息，请参阅 how-does-it-work。

支持的模型/模块

[!NOTE]

部分模型仅在 MMDet<3.0 上进行了测试。如果您发现任何无法成功转换的模型， 请在问题中报告。

• Faster R-CNN
• Cascade R-CNN
• Double-Head R-CNN
• Group Normalization
• Weight Standardization
• DCN
• SSD
• RetinaNet
• Libra R-CNN
• FCOS
• Fovea
• CARAFE
• FreeAnchor
• RepPoints
• NAS-FPN
• ATSS
• PAFPN
• FSAF
• GCNet
• Guided Anchoring
• Generalized Attention
• Dynamic R-CNN
• Hybrid Task Cascade
• DetectoRS
• Side-Aware Boundary Localization
• YOLOv3
• PAA
• CornerNet（开发中）
• Generalized Focal Loss
• Grid RCNN
• VFNet
• GROIE
• Mask R-CNN（实验）
• Cascade Mask R-CNN（实验）
• Cascade RPN
• DETR
• YOLOX

测试环境：

• torch=2.2.0
• tensorrt=8.6.1
• mmdetection=3.3.0
• cuda=11.7

常见问题解答

如果您遇到任何问题，请阅读 此页面。

许可证

本项目采用 Apache 2.0 许可证发布。

mmdetection-to-tensorrt 快速上手指南

本工具旨在将 MMDetection (>=3.0) 模型直接转换为 TensorRT 引擎，实现端到端的高效部署。它跳过 ONNX 中间格式，支持 FP16、动态输入形状及批量推理。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• GPU: NVIDIA GPU (支持 CUDA)
• 核心依赖版本参考:
  • PyTorch: 2.2.0
  • TensorRT: 8.6.1
  • CUDA: 11.7
  • MMDetection: 3.3.0

前置依赖安装

1. 安装 MMDetection 推荐使用 openmim 进行安装（国内用户可配置 pip 镜像源加速）：

   pip install openmim
   mim install mmdet==3.3.0
   

2. 安装 torch2trt_dynamic 该项目依赖特定的动态分支版本：

   git clone https://github.com/grimoire/torch2trt_dynamic.git torch2trt_dynamic
   cd torch2trt_dynamic
   pip install -e .
   

3. 安装 amirstan_plugin 需先安装官方 TensorRT，然后编译插件：

   # 克隆仓库
   git clone --depth=1 https://github.com/grimoire/amirstan_plugin.git
   cd amirstan_plugin
   
   # 初始化子模块
   git submodule update --init --progress --depth=1
   
   # 编译
   mkdir build
   cd build
   # 请将 ${TENSORRT_DIR} 替换为您实际的 TensorRT 安装路径
   cmake -DTENSORRT_DIR=${TENSORRT_DIR} ..
   make -j10
   

   重要：编译完成后，必须设置环境变量（建议写入 ~/.bashrc）：

   export AMIRSTAN_LIBRARY_PATH=${amirstan_plugin_root}/build/lib
   

安装步骤

您可以选择直接在宿主机安装或使用 Docker。

方式一：宿主机安装

git clone https://github.com/grimoire/mmdetection-to-tensorrt.git
cd mmdetection-to-tensorrt
pip install -e .

方式二：Docker 安装（推荐）

如果您希望避免复杂的环境配置，可以使用预构建的 Docker 镜像。

1. 构建镜像

   sudo docker build -t mmdet2trt_docker:v1.0 docker/
   

2. 运行容器 挂载数据目录并启动：

   sudo docker run --gpus all -it --rm -v ${your_data_path}:${bind_path} mmdet2trt_docker:v1.0
   

基本使用

方法 A：命令行工具 (CLI)

这是最简单的转换方式。准备好模型的配置文件 (config.py) 和权重文件 (checkpoint.pth)。

# 转换过程可能需要几分钟
mmdet2trt ${CONFIG_PATH} ${CHECKPOINT_PATH} ${OUTPUT_PATH}

示例（Docker 环境下）：

sudo docker run --gpus all -it --rm -v ${your_data_path}:${bind_path} mmdet2trt_docker:v1.0 ${bind_path}/config.py ${bind_path}/checkpoint.pth ${bind_path}/output.trt

方法 B：Python API

通过 Python 代码可以更灵活地控制转换参数（如动态形状、FP16 模式）。

# 定义动态输入形状范围
shape_ranges=dict(
    x=dict(
        min=[1,3,320,320],
        opt=[1,3,800,1344],
        max=[1,3,1344,1344],
    )
)

# 执行转换
trt_model = mmdet2trt(cfg_path,
                      weight_path,
                      shape_ranges=shape_ranges,
                      fp16_mode=True)

# 保存转换后的模型状态
torch.save(trt_model.state_dict(), save_model_path)

# 如果需要用于 C++ API，单独保存 engine
with open(save_engine_path, mode='wb') as f:
    f.write(trt_model.state_dict()['engine'])

推理测试

转换完成后，可以使用封装好的 Detector 进行推理，输出格式与原 MMDetection 保持一致。

from mmdet.apis import inference_detector
from mmdet2trt.apis import create_wrap_detector

# 创建包装后的检测器
trt_detector = create_wrap_detector(trt_model, cfg_path, device_id)

# 执行推理
result = inference_detector(trt_detector, image_path)

注意：

• 引擎的输入是预处理后的 Tensor。
• 引擎输出顺序为：num_dets, bboxes, scores, class_ids。若开启 enable_mask，还会输出 mask。
• 输出的 bboxes 未除以 scale_factor，使用时请注意坐标缩放。