Skip to content

TensorRT-YOLO: A high-performance, easy-to-use YOLO deployment toolkit for NVIDIA, powered by TensorRT plugins and CUDA Graph, supporting C and Python.

License

Notifications You must be signed in to change notification settings

laugh12321/TensorRT-YOLO

Repository files navigation

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 推理的支持,旨在提供📦开箱即用的部署体验。包括 目标检测实例分割图像分类姿态识别旋转目标检测视频分析等任务场景,满足开发者多场景的部署需求。

Detect Segment
Pose OBB

✨ 主要特性

  • 多样化的YOLO支持:全面兼容YOLOv3至YOLOv11以及PP-YOLOE和PP-YOLOE ,满足不同版本需求。
  • 多场景应用:提供DetectSegmentClassifyPoseOBB等多样化场景的示例代码。
  • 模型优化与推理加速
    • 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版本,确保与最新技术兼容。

🔮 文档教程

💨 快速开始

🔸 前置依赖

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

🎆 快速安装

Important

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

Python SDK快速开始

Python CLI 推理示例

Note

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

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

  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

DeployDetDeployOBBDeploySegDeployPose 分别对应于检测(Detect)、方向边界框(OBB)、分割(Segment)和姿态估计(Pose)模型。

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

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

# 初始化模型
model = DeployDet("yolo11n-with-plugin.engine")
# 加载图片
im = cv2.imread("test_image.jpg")
# 模型预测
result = model.predict(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
print(f"==> detect result: {result}")
# 可视化
labels = generate_labels_with_colors("labels.txt")
vis_im = visualize(im, result, labels)
cv2.imwrite("vis_image.jpg", vis_im)

C SDK快速开始

Note

DeployDetDeployOBBDeploySegDeployPose 分别对应于检测(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() {
    // 初始化模型
    auto model = deploy::DeployDet("yolo11n-with-plugin.engine");
    // 加载图片
    cv::Mat cvim = cv::imread("test_image.jpg");
    cv::cvtColor(cvim, cvim, cv::COLOR_BGR2RGB);
    deploy::Image im(cvim.data, cvim.cols, cvim.rows);
    // 模型预测
    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
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