YOLOv8-TensorRT

YOLOv8-TensorRT 项目介绍

YOLOv8-TensorRT项目是一个利用TensorRT技术加速YOLOv8的项目，旨在提升深度学习模型的推理速度和性能。

环境准备

要成功运行YOLOv8-TensorRT，用户需要：

1. 安装CUDA：
   • 推荐使用CUDA 11.4或更高版本。
   • 请按照NVIDIA的官方文档进行安装。
2. 安装TensorRT：
   • 推荐使用TensorRT 8.4或更高版本。
   • 请访问NVIDIA的官方网站获取TensorRT。
3. 安装Python依赖项：
   • 使用命令pip install -r requirements.txt安装必要的Python库。
4. 安装Ultralytics包：
   • 这一步是为ONNX导出或TensorRT API构建库需要。
   • 使用命令pip install ultralytics进行安装。
5. 准备PyTorch模型参数文件：
   • 例如，yolov8s.pt或yolov8s-seg.pt。

正常使用方法

如果用户从Ultralytics仓库获得ONNX模型，则需要自行构建引擎。同时，用户只能使用C++推理代码来反序列化引擎并进行推理。

导出含NMS的End2End 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 opset版本，默认为11。
• --sim：是否简化ONNX模型。
• --input-shape：模型的输入形状，应该是四维的。
• --device：导出引擎所使用的CUDA设备。

从ONNX构建End2End引擎

1. 使用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

2. 使用Trtexec工具导出引擎

可以通过trtexec工具导出TensorRT引擎。

示例命令：

/usr/src/tensorrt/bin/trtexec --onnx=yolov8s.onnx --saveEngine=yolov8s.engine --fp16

推理

1. 使用Python脚本进行推理

用户可以通过infer-det.py脚本对图像进行推理。

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

2. 使用C++进行推理

在路径csrc/detect/end2end下使用C++进行推理。

构建步骤及样例使用命令：

./yolov8 yolov8s.engine data/bus.jpg

其他部署选项

• 支持使用TensorRT Segment、Pose、Cls、Obb等多种部署。
• 提供在Jetson平台上的测试。

免责声明

该项目明确指出不支持YOLOv8-seg模型，建议用户使用其他选项进行相关操作。此外，如果需要抛弃PyTorch而完全使用TensorRT进行推理，性能可能会有所下降。为高效推理，可考虑使用cuda-python或pycuda。

以上概述了YOLOv8-TensorRT项目的要点，旨在帮助用户更快地理解和部署该项目。