BEVFormer_tensorrt: 基于TensorRT的高效BEV 3D目标检测部署方案

BEVFormer_tensorrt:让BEV 3D目标检测更快更高效

在自动驾驶和机器人技术快速发展的今天,3D目标检测作为一项核心技术正受到越来越多的关注。然而,3D目标检测模型通常计算量巨大,难以在实际应用中实现实时处理。为了解决这一问题,研究人员开发了BEVFormer_tensorrt项目,旨在通过TensorRT加速和优化BEV(Bird's Eye View)3D目标检测模型的推理过程。

项目概述

BEVFormer_tensorrt是一个专注于将BEV 3D目标检测模型(如BEVFormer、BEVDet等)部署到TensorRT平台上的开源项目。该项目不仅支持FP32/FP16/INT8等多种精度的推理,还实现了一系列自定义TensorRT插件,以进一步提升模型的推理速度和内存效率。

主要特性

1. 多精度支持:项目支持FP32、FP16和INT8三种精度的推理,可以根据实际需求在精度和速度之间进行权衡。

2. 自定义TensorRT插件:为了优化BEVFormer在TensorRT上的推理性能,项目实现了多个支持nv_half、nv_half2和int8的TensorRT插件。这些插件包括:

   • Grid Sampler
   • Multi-scale Deformable Attention
   • Modulated Deformable Conv2d
   • Rotate
   • Inverse
   • BEV Pool V2
   • Flash Multi-Head Attention

3. 显著的性能提升:通过使用自定义插件和INT8量化,BEVFormer base模型的推理速度提升了5倍以上,同时模型大小减少了90%以上,GPU内存使用降低了80%以上。

4. 广泛的模型支持:除了BEVFormer,项目还支持BEVDet等其他BEV 3D目标检测模型,以及MMDetection中的常见2D目标检测模型(如YOLOx和CenterNet)。

性能对比

为了直观展示BEVFormer_tensorrt的优势,我们来看一下BEVFormer base模型在不同配置下的性能对比:

配置	NDS/mAP	FPS	模型大小(MB)	GPU内存(MB)
PyTorch原始版本	0.517/0.416	2.4	265	5435
TensorRT(FP32)	0.517/0.416	3.0	292	5715
TensorRT(FP16)	0.517/0.416	4.9	148	3417
TensorRT(FP16/INT8)	0.514/0.413	8.0	131	2429

从表中可以看出,使用TensorRT和自定义插件后,BEVFormer base模型在保持准确率基本不变的情况下,推理速度从2.4 FPS提升到了8.0 FPS,提升了233%。同时,模型大小从265MB减少到131MB,GPU内存使用从5435MB降低到2429MB。这些改进对于在实际应用中部署BEV 3D目标检测模型具有重要意义。

安装与使用

BEVFormer_tensorrt的安装过程相对复杂,需要配置CUDA、cuDNN、TensorRT等环境。项目提供了两种安装方式:Docker和从源码安装。对于大多数用户来说,使用Docker是更简单和推荐的方式。

以下是使用Docker安装的基本步骤:

1. 克隆项目仓库:

git clone git@github.com:DerryHub/BEVFormer_tensorrt.git
cd BEVFormer_tensorrt

2. 构建Docker镜像:

docker build -t trt85 -f docker/Dockerfile .

3. 运行Docker容器:

docker run -it --gpus all -v ${PROJECT_DIR}:/workspace/BEVFormer_tensorrt/ \
-v /path/to/can_bus:/workspace/BEVFormer_tensorrt/data/can_bus \
-v /path/to/coco:/workspace/BEVFormer_tensorrt/data/coco \
-v /path/to/nuscenes:/workspace/BEVFormer_tensorrt/data/nuscenes \
--shm-size 8G trt85 /bin/bash

4. 在容器内编译和安装自定义TensorRT插件:

cd /workspace/BEVFormer_tensorrt/TensorRT/build
cmake .. -DCMAKE_TENSORRT_PATH=/usr
make -j$(nproc)
make install

模型转换与评估

BEVFormer_tensorrt提供了一系列脚本,用于将PyTorch模型转换为ONNX格式,然后再转换为TensorRT引擎。以BEVFormer base模型为例,转换和评估的步骤如下:

1. 将.pth模型转换为.onnx格式:

sh samples/bevformer/plugin/base/pth2onnx.sh -d ${gpu_id}

2. 将.onnx模型转换为TensorRT引擎(以FP16精度为例):

sh samples/bevformer/plugin/base/onnx2trt_fp16.sh -d ${gpu_id}

3. 使用TensorRT引擎进行评估:

sh samples/bevformer/plugin/base/trt_evaluate_fp16.sh -d ${gpu_id}

通过这些步骤,用户可以轻松地将BEVFormer模型部署到TensorRT上,并验证其性能和准确率。

结论与展望

BEVFormer_tensorrt项目为BEV 3D目标检测模型的实际应用提供了一个高效的部署方案。通过利用TensorRT的硬件加速能力和自定义优化插件,该项目显著提高了模型的推理速度,同时大幅降低了内存占用和模型大小。这些改进使得在资源受限的嵌入式设备上部署复杂的3D目标检测模型成为可能。

未来,BEVFormer_tensorrt项目可能会继续优化现有插件的性能,并支持更多的BEV 3D目标检测模型。同时,随着硬件技术的进步和新的优化技术的出现,我们可以期待BEV 3D目标检测在实时性和资源效率方面取得进一步的突破。

对于自动驾驶、机器人导航等领域的开发者和研究人员来说,BEVFormer_tensorrt无疑是一个值得关注和尝试的项目。它不仅提供了现成的部署方案,也为如何优化大型视觉模型以适应实际应用需求提供了宝贵的经验和思路。

总的来说,BEVFormer_tensorrt代表了计算机视觉领域中模型优化和部署的一个重要方向。随着3D感知技术在各个领域的广泛应用,类似的优化工作将变得越来越重要,为推动自动驾驶和机器人技术的发展做出重要贡献。