TensorRT-YOLO 使用教程

TensorRT-YOLO 使用教程

1. 项目介绍

TensorRT-YOLO 是一个支持 YOLOv3、YOLOv5、YOLOv6、YOLOv7、YOLOv8、YOLOv9、YOLOv10、PP-YOLOE 和 PP-YOLOE+ 的推理加速项目。该项目使用 NVIDIA TensorRT 进行优化，不仅集成了 EfficientNMS TensorRT 插件以增强后处理效果，还使用了 CUDA 核函数以及 CUDA 图来加速推理。TensorRT-YOLO 提供了 C++ 和 Python 推理的支持，旨在提供快速而优化的目标检测解决方案。

2. 项目快速启动

环境要求

• CUDA 版本 >= 11.6
• TensorRT 版本 >= 8.6

快速编译安装

1. 克隆项目仓库：

   git clone https://github.com/laugh12321/TensorRT-YOLO.git
   cd TensorRT-YOLO
   

2. 安装依赖：

   pip install -r requirements.txt
   

3. 编译项目：

   xmake build
   

使用 CLI 模型导出

1. 导出 ONNX 模型：

   python export.py --weights yolov5s.pt --include onnx
   

2. 导出 TensorRT 模型：

   python export.py --weights yolov5s.onnx --include trt
   

模型推理示例

1. 使用 Python 进行推理：

   from tensorrt_yolo import TensorRTYOLO
   
   model = TensorRTYOLO('yolov5s.trt')
   image = model.infer('test.jpg')
   image.save('output.jpg')
   

2. 使用 C++ 进行推理：

   #include "tensorrt_yolo.h"
   
   int main() {
       TensorRTYOLO model("yolov5s.trt");
       cv::Mat image = cv::imread("test.jpg");
       cv::Mat output = model.infer(image);
       cv::imwrite("output.jpg", output);
       return 0;
   }
   

3. 应用案例和最佳实践

视频分析示例

TensorRT-YOLO 可以用于实时视频流的目标检测。以下是一个简单的视频分析示例：

import cv2
from tensorrt_yolo import TensorRTYOLO

model = TensorRTYOLO('yolov5s.trt')
cap = cv2.VideoCapture('test.mp4')

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    output = model.infer(frame)
    cv2.imshow('TensorRT-YOLO', output)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

最佳实践

• 模型优化：使用 TensorRT 的 INT8 量化功能进一步优化模型推理速度。
• 多线程推理：在多核 CPU 环境下，使用多线程并行推理以提高吞吐量。
• CUDA 图加速：利用 CUDA 图技术减少 GPU 启动延迟，提升推理性能。

4. 典型生态项目

1. NVIDIA TensorRT

TensorRT 是 NVIDIA 提供的高性能深度学习推理库，支持多种深度学习框架的模型优化和加速。TensorRT-YOLO 充分利用了 TensorRT 的优化功能，提供了高效的 YOLO 模型推理。

2. YOLO 系列

YOLO（You Only Look Once）是一系列实时目标检测算法，广泛应用于自动驾驶、视频监控等领域。TensorRT-YOLO 支持 YOLOv3 到 YOLOv10 的多个版本，为用户提供了丰富的选择。

3. EfficientNMS TensorRT 插件

EfficientNMS 是 TensorRT 的一个插件，用于加速目标检测中的非极大值抑制（NMS）操作。TensorRT-YOLO 集成了该插件，显著提升了后处理速度。

4. CUDA 核函数

CUDA 核函数是 NVIDIA 提供的并行计算 API，用于在 GPU 上执行高性能计算任务。TensorRT-YOLO 使用 CUDA 核函数加速前处理和后处理步骤，进一步提升了推理速度。

通过以上模块的介绍和示例，您可以快速上手 TensorRT-YOLO 项目，并了解其在实际应用中的最佳实践和生态项目。