YOLOv8-TensorRT: 使用TensorRT加速YOLOv8目标检测

YOLOv8-TensorRT

YOLOv8-TensorRT: 使用TensorRT加速YOLOv8目标检测

YOLOv8是目前最先进的目标检测算法之一,具有出色的精度和速度。而TensorRT作为NVIDIA推出的高性能深度学习推理优化器和运行时环境,能够显著提升模型的推理速度。本文将详细介绍如何使用TensorRT加速YOLOv8,助力开发者实现高性能的YOLOv8部署。

环境准备

在开始使用YOLOv8-TensorRT之前,我们需要做好以下环境准备:

安装CUDA

建议安装CUDA 11.4或更高版本。可以从NVIDIA CUDA官网下载并安装。
安装TensorRT

推荐安装TensorRT 8.4或更高版本。可以从NVIDIA TensorRT官网下载并安装。
安装Python依赖

执行以下命令安装所需的Python包:
```
pip install -r requirements.txt
```
安装ultralytics包

这个包用于ONNX导出或TensorRT API构建:
```
pip install ultralytics
```
准备PyTorch权重文件

例如yolov8s.pt或yolov8s-seg.pt。

🚀 注意: 请使用最新版本的CUDA和TensorRT以获得最佳性能。如果必须使用较低版本,请仔细阅读相关问题。

导出端到端ONNX模型

我们可以使用ultralytics API导出ONNX模型,并同时将后处理(如边界框解码和NMS)添加到ONNX模型中:

python3 export-det.py \
--weights yolov8s.pt \
--iou-thres 0.65 \
--conf-thres 0.25 \
--topk 100 \
--opset 11 \
--sim \
--input-shape 1 3 640 640 \
--device cuda:0

参数说明:

--weights: 训练好的PyTorch模型
--iou-thres: NMS插件的IOU阈值
--conf-thres: NMS插件的置信度阈值
--topk: 检测框的最大数量
--opset: ONNX操作集版本,默认为11
--sim: 是否简化ONNX模型
--input-shape: 模型输入shape,应为4维
--device: 导出引擎的CUDA设备

执行后将得到一个与输入权重同名的ONNX模型。

构建TensorRT引擎

有两种方法可以从ONNX模型构建TensorRT引擎:

使用TensorRT ONNX Python API

使用build.py脚本从ONNX导出TensorRT引擎:
```
python3 build.py \
--weights yolov8s.onnx \
--iou-thres 0.65 \
--conf-thres 0.25 \
--topk 100 \
--fp16  \
--device cuda:0
```
参数说明:
- --weights: 下载的ONNX模型
- --iou-thres: NMS插件的IOU阈值
- --conf-thres: NMS插件的置信度阈值
- --topk: 检测框的最大数量
- --fp16: 是否导出半精度引擎
- --device: 导出引擎的CUDA设备
使用Trtexec工具

也可以使用NVIDIA提供的trtexec工具导出TensorRT引擎:
```
/usr/src/tensorrt/bin/trtexec \
--onnx=yolov8s.onnx \
--saveEngine=yolov8s.engine \
--fp16
```
注意:trtexec的安装路径可能因TensorRT的安装方式而异。

推理部署

Python推理

可以使用infer-det.py脚本进行图像推理:

python3 infer-det.py \
--engine yolov8s.engine \
--imgs data \
--show \
--out-dir outputs \
--device cuda:0

参数说明:

--engine: 导出的引擎文件
--imgs: 要检测的图像路径
--show: 是否显示检测结果
--out-dir: 保存检测结果图像的位置
--device: 使用的CUDA设备
--profile: 是否对TensorRT引擎进行性能分析

C++推理

在csrc/detect/end2end目录下提供了C++推理代码。使用前需要在CMakeLists.txt中设置自己的库路径,并在main.cpp中修改CLASS_NAMES和COLORS。

构建步骤:

export root=${PWD}
cd csrc/detect/end2end
mkdir -p build && cd build
cmake ..
make
mv yolov8 ${root}
cd ${root}

使用方法:

# 推理单张图片
./yolov8 yolov8s.engine data/bus.jpg

# 推理多张图片
./yolov8 yolov8s.engine data

# 推理视频
./yolov8 yolov8s.engine data/test.mp4

其他功能

TensorRT分割模型部署

详情请参考Segment.md。
TensorRT姿态估计模型部署

详情请参考Pose.md。
TensorRT分类模型部署

详情请参考Cls.md。
DeepStream检测模型部署

详情请参考csrc/deepstream/README.md。
Jetson部署

仅在Jetson-NX 4GB上测试过,详情请参考Jetson.md。
引擎性能分析

如果想对TensorRT引擎进行性能分析:
```
python3 trt-profile.py --engine yolov8s.engine --device cuda:0
```
无PyTorch推理

如果需要脱离PyTorch使用TensorRT推理,可以参考infer-det-without-torch.py。该脚本支持使用cuda-python或pycuda进行推理,但性能可能不如PyTorch版本。

安装依赖:
```
pip install cuda-python
# 或
pip install pycuda
```
使用方法:
```
python3 infer-det-without-torch.py \
--engine yolov8s.engine \
--imgs data \
--show \
--out-dir outputs \
--method cudart
```
参数--method可选cudart或pycuda,默认为cudart。

通过本文的详细介绍,相信读者已经对如何使用TensorRT加速YOLOv8有了全面的了解。YOLOv8-TensorRT不仅提供了完整的模型转换和部署流程,还支持多种任务类型和部署平台,是一个功能强大且灵活的工具。希望这个项目能够帮助更多开发者实现高性能的YOLOv8部署,推动计算机视觉技术的应用和发展。

YOLOv8-TensorRT架构图

YOLOv8-TensorRT项目为开发者提供了一个强大的工具,使其能够轻松地将YOLOv8模型与TensorRT结合,实现高效的目标检测。无论是在边缘设备还是在服务器上,这个项目都能帮助开发者充分发挥硬件性能,加速AI应用的部署。我们期待看到更多基于YOLOv8-TensorRT的创新应用,为计算机视觉领域带来新的可能性。