English | 简体中文

🚀 TensorRT YOLO

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

Detect	Segment
Pose	OBB

✨ 主要特性

• 多样化的YOLO支持：全面兼容YOLOv3至YOLOv11以及PP-YOLOE和PP-YOLOE+，满足不同版本需求。
• 多场景应用：提供Detect、Segment、Pose、OBB等多样化场景的示例代码。
• 模型优化与推理加速：
  • ONNX支持：支持ONNX模型的静态和动态导出，包括TensorRT自定义插件支持，简化模型部署流程。
  • TensorRT集成：集成TensorRT插件，包括自定义插件，加速Detect, Segment, Pose, OBB等场景的后处理，提升推理效率。
  • CUDA加速：利用CUDA核函数优化前处理，CUDA图技术加速推理流程，实现高性能计算。
• 语言支持：支持C++和Python（通过Pybind11映射，提升Python推理速度），满足不同编程语言需求。
• 部署便捷性：
  • 动态库编译：提供动态库编译支持，方便调用和部署。
  • 无第三方依赖：全部功能使用标准库实现，无需额外依赖，简化部署流程。
• 快速开发与部署：
  • CLI工具：提供命令行界面（CLI）工具，实现快速模型导出和推理。
  • 跨平台支持：支持Windows、Linux、ARM、x86等多种设备，适应不同硬件环境。
  • Docker部署：支持Docker一键部署，简化环境配置和部署流程。
• TensorRT兼容性：兼容TensorRT 10.x版本，确保与最新技术兼容。

🔮 文档教程

• 安装文档
  • 📦 快速编译安装
• 快速开始
  • ✴️ Python SDK快速使用
  • ✴️ C++ SDK快速使用
• 使用示例
  • 目标检测 示例
  • 实例分割 示例
  • 姿态识别 示例
  • 旋转目标检测 示例
  • 📹视频分析 示例
• API文档
  • Python API文档（⚠️ 未实现）
  • C++ API文档（⚠️ 未实现）
• 常见问题
  • ⚠️ 收集中 ...
• 模型支持列表
  • 🖥️ 模型支持列表

💨 快速开始

🔸 前置依赖

• 推荐 CUDA 版本 >= 11.6
• 推荐 TensorRT 版本 >= 8.6.1 （TensorRT 最低版本 8.6.1）
• OS: Linux x86_64 (推荐) arm / Windows /

🎆 快速安装

• 参考📦 快速编译安装文档

Python SDK快速开始

Python CLI 推理示例

Note

使用 --cudaGraph 选项可以显著提升推理速度，但需知此功能仅适用于静态模型。

通过 -m, --mode 参数可以选择不同的模型类型，其中 0 代表检测（Detect）、1 代表旋转边界框（OBB）、2 代表分割（Segment）、3 代表姿态估计（Pose）。

1. 使用 tensorrt_yolo 库的 trtyolo 命令行工具进行推理。运行以下命令查看帮助信息：

   trtyolo infer --help

2. 运行以下命令进行推理：

   trtyolo infer -e models/yolo11n.engine -m 0 -i images -o output -l labels.txt --cudaGraph

   推理结果将保存至 output 文件夹，并生成可视化结果。

Python 推理示例

Note

DeployDet、DeployOBB、DeploySeg 和 DeployPose 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）和姿态估计（Pose）模型。

对于这些模型，CG 版本利用 CUDA Graph 来进一步加速推理过程，但请注意，这一功能仅限于静态模型。

import cv2
from tensorrt_yolo.infer import DeployCGDet, DeployDet, generate_labels_with_colors, visualize

use_cudaGraph = True
engine_path = "yolo11n-with-plugin.engine"
model = DeployCGDet(engine_path) if use_cudaGraph else DeployDet(engine_path)

im = cv2.imread("test_image.jpg")
result = model.predict(cv2.cvtColor(im, cv2.COLOR_BGR2RGB)) # model 接收的图片必须是RGB格式
print(f"==> detect result: {result}")

# 可视化
labels = generate_labels_with_colors("labels.txt")
vis_im = visualize(image, result, labels)
cv2.imwrite("vis_image.jpg", vis_im)

C++ SDK快速开始

Important

在进行推理之前，请参考🔧 CLI 导出模型文档，导出适用于该项目推理的ONNX模型并构建为TensorRT引擎。

Note

DeployDet、DeployOBB、DeploySeg 和 DeployPose 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）和姿态估计（Pose）模型。

对于这些模型，CG 版本利用 CUDA Graph 来进一步加速推理过程，但请注意，这一功能仅限于静态模型。

#include <opencv2/opencv.hpp>
// 为了方便调用，模块除使用 CUDA、TensorRT 其余均使用标准库实现
#include "deploy/vision/inference.hpp"
#include "deploy/vision/result.hpp"

int main(int argc, char* argv[]) {
    bool useCudaGraph = true;
    deploy::DeployBase model;
    if (useCudaGraph) {
        model = deploy::DeployCGDet("yolo11n-with-plugin.engine");
    } else {
        model = deploy::DeployDet("yolo11n-with-plugin.engine");
    }
    auto cvim = cv::imread("test_image.jpg");

    cv::cvtColor(cvim, cvim, cv::COLOR_BGR2RGB);
    deploy::Image im(cvim.data, cvim.cols, cvim.rows); // model 接收的图片必须是RGB格式
    deploy::DetResult result = 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+	✅	✅
OBB	ultralytics/ultralytics	✅	✅
Pose	ultralytics/ultralytics	✅	✅
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	✅	✅

☕ 请作者喝杯咖啡

开源不易，如果本项目有帮助到你的话，可以考虑请作者喝杯咖啡，你的支持是开发者持续维护的最大动力~

📄 许可证

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

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

📞 联系方式

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