laugh12321/TensorRT-YOLO

TensorRT‑YOLO – 在 NVIDIA 硬件上实现高性能 YOLO 推理

是什么 – 一个 C++/Python 库，封装了 TensorRT 以在 NVIDIA GPU 和 Jetson 设备上运行整个 YOLO 系列（v3‑v26, YOLO‑World, YOLO‑E, YOLO‑Master 等）。它附带了自定义的 TensorRT 插件、CUDA 内核和 CUDA‑Graph 支持，旨在使端到端流水线（预处理 → 推理 → 后处理）尽可能快，同时保持 API 的简洁性。

核心能力

• 广泛的模型覆盖 – 支持许多 YOLO 变体的检测、分割、分类、姿态估计和旋转框（OBB）模型。
• 性能优化技巧 – CUDA 加速的预处理、用于后处理的 TensorRT 插件、多上下文并行推理、Jetson 上的 Zero‑Copy 以及可选的 CUDA‑Graph 执行。
• 易于集成 – 提供单头文件 C++ 接口 (trtyolo.hpp) 和 Python 绑定 (trtyolo)，两者都无需额外的第三方库即可使用。
• 跨平台 – 支持 Linux (x86_64 & ARM) 和 Windows；提供 Docker 镜像以实现一键式设置。
• 零依赖构建 – C++ 库可以编译为仅头文件模块；Python wheel 包集成了绑定，无需用户手动链接 CUDA/TensorRT。
• CLI & 示例 – 提供针对检测、分割、分类、姿态、OBB 和视频流水线的即用型演示程序。

典型工作流程

1. 转换 – 使用配套的 trtyolo-export 工具将导出的 ONNX YOLO 模型转换为 TensorRT‑YOLO 引擎格式。
2. 构建 – 构建库（CMake） – 如果需要 Python API，请启用 BUILD_PYTHON=ON。
3. 加载 引擎于代码中：

   from trtyolo import TRTYOLO
   model = TRTYOLO('yolo11n-with-plugin.engine', task='detect', profile=True, swap_rb=True)
   result = model.predict(cv2.imread('img.jpg'))
   

   或者在 C++ 中：

   auto det = std::make_unique<trtyolo::DetectModel>("yolo11n-with-plugin.engine", option);
   auto out = det->predict(image);
   

4. 可选 – 调用 model.clone()（或 det->clone()）以获取用于多线程服务器的线程安全副本。
5. 性能分析 – model.profile() 返回吞吐量和延迟指标。

为什么重要 – YOLO 模型在实时视觉任务中非常流行，但在 NVIDIA 硬件上获得最佳速度通常需要手工编写 TensorRT 插件并进行精细的 CUDA 调优。TensorRT‑YOLO 封装了这些优化并提供了简洁的 API，让开发者能够专注于应用逻辑，而不是底层的引擎架构。

许可证 & 社区 – GPL‑3.0，拥有活跃的问题跟踪、Issue tracking、Bilibili 教程系列以及用于捐赠的赞助页面。

---

快速开始 (Linux)

git clone https://github.com/laugh12321/TensorRT-YOLO && cd TensorRT-YOLO
pip install "pybind11[global]"
cmake -S . -B build -DTRT_PATH=/path/to/TensorRT -DBUILD_PYTHON=ON -DCMAKE_INSTALL_PREFIX=$HOME/trtyolo
cmake --build build -j$(nproc) --target install
# Build Python wheel
pip install --upgrade build
python -m build --wheel && pip install dist/trtyolo-*.whl

现在你可以运行提供的 examples/detect 脚本或按照上述方式编写自己的代码。