Skip to content

Latest commit

 

History

History
346 lines (291 loc) · 14.2 KB

README.md

File metadata and controls

346 lines (291 loc) · 14.2 KB

English | 简体中文

🚀 TensorRT YOLO

GitHub License GitHub Release GitHub commit activity GitHub Repo stars GitHub forks

🚀 TensorRT-YOLO 是一款专为 NVIDIA 设备设计的易用灵活、极致高效的YOLO系列推理部署工具。项目不仅集成了 TensorRT 插件以增强后处理效果，还使用了 CUDA 核函数以及 CUDA 图来加速推理。TensorRT-YOLO 提供了 C++ 和 Python 推理的支持，旨在提供📦开箱即用的部署体验。包括 目标检测、实例分割、图像分类、姿态识别、旋转目标检测、视频分析等任务场景，满足开发者多场景的部署需求。

✨ 主要特性

🎯 多样化的 YOLO 支持

  • 全面兼容：支持 YOLOv3 至 YOLOv11 全系列模型，以及 PP-YOLOE 和 PP-YOLOE+，满足多样化需求。
  • 灵活切换：提供简洁易用的接口，支持不同版本 YOLO 模型的快速切换。
  • 多场景应用：提供丰富的示例代码，涵盖Detect、Segment、Classify、Pose、OBB等多种应用场景。

🚀 性能优化

  • CUDA 加速：通过 CUDA 核函数优化前处理流程，并采用 CUDA 图技术加速推理过程。
  • TensorRT 集成：深度集成 TensorRT 插件，显著加速后处理，提升整体推理效率。
  • 多 Context 推理：支持多 Context 并行推理，最大化硬件资源利用率。
  • 显存管理优化：适配多架构显存优化策略（如 Jetson 的 Zero Copy 模式），提升显存效率。

🛠️ 易用性

  • 开箱即用：提供全面的 C++ 和 Python 推理支持，满足不同开发者需求。
  • CLI 工具：内置命令行工具，支持快速模型导出与推理，提升开发效率。
  • Docker 支持：提供 Docker 一键部署方案，简化环境配置与部署流程。
  • 无第三方依赖：全部功能使用标准库实现，无需额外依赖，简化部署流程。
  • 部署便捷：提供动态库编译支持，方便调用和部署。

🌐 兼容性

  • 多平台支持：全面兼容 Windows、Linux、ARM、x86 等多种操作系统与硬件平台。
  • TensorRT 兼容：完美适配 TensorRT 10.x 版本，确保与最新技术生态无缝衔接。

🔧 灵活配置

  • 预处理参数自定义：支持多种预处理参数灵活配置，包括 通道交换 (SwapRB)、归一化参数、边界值填充。

🔮 文档教程

💨 快速开始

1. 前置依赖

  • CUDA：推荐版本 ≥ 11.0.1
  • TensorRT：推荐版本 ≥ 8.6.1
  • 操作系统：Linux (x86_64 或 arm)（推荐）；Windows 亦可支持

2. 安装

3. 模型导出

  • 参考 🔧 模型导出 文档，导出适用于该项目推理的ONNX模型并构建为TensorRT引擎。

4. 推理示例

Note

ClassifyModel、DetectModel、OBBModel、SegmentModel 和 PoseModel 分别对应于图像分类（Classify）、检测（Detect）、方向边界框（OBB）、分割（Segment）、姿态估计（Pose）和模型。

  • 使用 Python 进行推理：

    import cv2
from tensorrt_yolo.infer import InferOption, DetectModel, generate_labels, visualize

# 配置推理选项
option = InferOption()
option.enable_swap_rb()

# 初始化模型
model = DetectModel("yolo11n-with-plugin.engine", option)

# 加载图片
im = cv2.imread("test_image.jpg")

# 模型预测
result = model.predict(im)
print(f"==> detect result: {result}")

# 可视化检测结果
labels = generate_labels("labels.txt")
vis_im = visualize(im, result, labels)
cv2.imwrite("vis_image.jpg", vis_im)

# 克隆模型并进行预测
clone_model = model.clone()
clone_result = clone_model.predict(im)
print(f"==> detect clone result: {clone_result}")

  • 使用 C++ 进行推理：

    #include <memory>
#include <opencv2/opencv.hpp>

// 为了方便调用，模块除使用CUDA、TensorRT外，其余均使用标准库实现
#include "deploy/model.hpp"  // 包含模型推理相关的类定义
#include "deploy/option.hpp"  // 包含推理选项的配置类定义
#include "deploy/result.hpp"  // 包含推理结果的定义

int main() {
    // 配置推理选项
    deploy::InferOption option;
    option.enableSwapRB();  // 启用通道交换（从BGR到RGB）

    // 初始化模型
    auto model = std::make_unique<deploy::DetectModel>("yolo11n-with-plugin.engine", option);

    // 加载图片
    cv::Mat cvim = cv::imread("test_image.jpg");
    deploy::Image im(cvim.data, cvim.cols, cvim.rows);

    // 模型预测
    deploy::DetResult result = model->predict(im);

    // 可视化（代码省略）
    // ...  // 可视化部分代码未提供，可根据需要实现

    // 克隆模型并进行预测
    auto clone_model = model->clone();
    deploy::DetResult clone_result = clone_model->predict(im);

    return 0;  // 程序正常结束
}

更多部署案例请参考模型部署示例 .

🖥️ 模型支持列表

Detect Segment
Pose OBB

符号说明: (1) ✅ : 已经支持; (2) ❔: 正在进行中; (3) ❎ : 暂不支持; (4) 🟢 : 导出自行实现，即可推理.

任务场景 模型 CLI 导出 推理部署
Detect ultralytics/yolov3 ✅ ✅
Detect ultralytics/yolov5 ✅ ✅
Detect meituan/YOLOv6 ❎ 参考官方导出教程 ✅
Detect WongKinYiu/yolov7 ❎ 参考官方导出教程 ✅
Detect WongKinYiu/yolov9 ❎ 参考官方导出教程 ✅
Detect THU-MIG/yolov10 ✅ ✅
Detect ultralytics/ultralytics ✅ ✅
Detect PaddleDetection/PP-YOLOE+ ✅ ✅
Segment ultralytics/yolov3 ✅ ✅
Segment ultralytics/yolov5 ✅ ✅
Segment meituan/YOLOv6-seg ❎ 参考tensorrt_yolo/export/head.py 自行实现 🟢
Segment WongKinYiu/yolov7 ❎ 参考tensorrt_yolo/export/head.py 自行实现 🟢
Segment WongKinYiu/yolov9 ❎ 参考tensorrt_yolo/export/head.py 自行实现 🟢
Segment ultralytics/ultralytics ✅ ✅
Classify ultralytics/yolov3 ✅ ✅
Classify ultralytics/yolov5 ✅ ✅
Classify ultralytics/ultralytics ✅ ✅
Pose ultralytics/ultralytics ✅ ✅
OBB ultralytics/ultralytics ✅ ✅

📄 许可证

TensorRT-YOLO采用 GPL-3.0许可证，这个OSI 批准的开源许可证非常适合学生和爱好者，可以推动开放的协作和知识分享。请查看LICENSE 文件以了解更多细节。

感谢您选择使用 TensorRT-YOLO，我们鼓励开放的协作和知识分享，同时也希望您遵守开源许可的相关规定。

📞 联系方式

对于 TensorRT-YOLO 的错误报告和功能请求，请访问 GitHub Issues！

🙏 致谢

Featured｜HelloGitHub