RT-DETR-V2 TensorRT C++ 部署

山水无移2024-10-29 23:20

RT-DETRv2_TensorRT_Cplusplus

前几天看到RT-DETR都到V3了，有同事问搞过DETR没，之前还真部署玩过，写了两篇博客【DETR tensorRT 部署】和【DETR tensorRT部署去除推理过程无用辅助头+fp16部署再次加速+解决转tensorrt 输出全为0问题的新方法】，当时部署DETR时有两个问题：第一个，不用预训练权重训练自己数据集MAP为0；第二个部署tensorRTs时输出全为0。第一个问题参考官方提供的加载预训练权重很容易解决，第二个问题部署tensorRT输出全为0的问题，网上资料很少，困扰了很久，反反复复思考，是找到了一个解决方法，但后续又详细研究了一下detr，最终想到一个解决的输出全为0的新方法，不仅部署简单而且速度更快。趁着这两天有时间抓紧折腾一下rt-detr，本来想部署一下rt-detr-v3，奈何代码还没开源，那就折腾一下rt-detr-v2。

RT-DETRv2 tensorrt C++ 部署

【完整代码】

本示例中，包含完整的代码、模型、测试图片、测试结果。

TensorRT版本：TensorRT-8.6.1.6

rt-detrv2 训练

训练参考官方开源代码。

导出onnx模型

在官方导出onnx 的基础上进行简单的调整，这里不需要动态batch，也不需要进行解码到输入分辨率，进行了如下调整：

python 复制代码

        def forward(self, images, orig_target_sizes):
            outputs = self.model(images)
            # outputs = self.postprocessor(outputs, orig_target_sizes)
            outputs = torch.sigmoid(outputs['pred_logits']), outputs['pred_boxes']
            return outputs


    model = Model()
    data = torch.rand(1, 3, 640, 640)

    torch.onnx.export(
        model, 
        data,
        args.output_file,
        input_names=['images'],
        output_names=['output1', 'output2'],
        opset_version=16, 
        verbose=False,
        do_constant_folding=True,
    )

最终导出的onnx结构如下：

onnx 运行结果

onnx 测试脚本【链接】

TensorRT C++ 部署

tensorrt 环境搭建参考官方文档，主要版本和cuda匹配。

1、修改Tensorrt使用版本

2、修改代码中模型对应的路径

3、编译运行

shellpower 复制代码

# 编译
cd RT-DETRv2_TensorRT_Cplusplus
mkdir build
cd build
cmake ..
make


# 运行
./detr_trt

tensorrt 运行结果

特别说明：本示例用fp16精度掉的非常多，默认使用的fp32。

运行时耗

本示例使用的是 rtdetrv2_r18vd_120e_coco.yml 模型，模型输入分别率640x640，显卡rtx4090，cuda12.5，fp32。

待解决问题

用FP16精度丢的很严重，目前还不知道为啥。