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检测模型部署
-
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结合,实现高效的目标检测。无论是在边缘设备还是在服务器上,这个项目都能帮助开发者充分发挥硬件性能,加速AI应用的部署。我们期待看到更多基于YOLOv8-TensorRT的创新应用,为计算机视觉领域带来新的可能性。
