laugh12321/TensorRT-YOLO

TensorRT‑YOLO – High‑performance YOLO inference on NVIDIA hardware

What it is – A C++/Python library that wraps TensorRT to run the whole YOLO family (v3‑v26, YOLO‑World, YOLO‑E, YOLO‑Master, etc.) on NVIDIA GPUs and Jetson devices. It ships with custom TensorRT plugins, CUDA kernels and CUDA‑Graph support to make the end‑to‑end pipeline (pre‑process → inference → post‑process) as fast as possible while keeping the API simple.

Key capabilities

• Broad model coverage – Detect, segment, classify, pose‑estimate and oriented‑bbox (OBB) models from many YOLO variants.
• Performance tricks – CUDA‑accelerated pre‑processing, TensorRT plugins for post‑processing, multi‑context parallel inference, Zero‑Copy on Jetson, and optional CUDA‑Graph execution.
• Easy integration – Single‑header C++ interface (trtyolo.hpp) and a Python binding (trtyolo), both usable without extra third‑party libraries.
• Cross‑platform – Works on Linux (x86_64 & ARM) and Windows; Docker image provided for one‑click setup.
• Zero‑dependency build – The C++ library can be compiled as a header‑only module; the Python wheel bundles the bindings without requiring the user to link CUDA/TensorRT manually.
• CLI & examples – Ready‑to‑run demos for detection, segmentation, classification, pose, OBB and video pipelines.

Typical workflow

1. Convert an ONNX‑exported YOLO model with the companion trtyolo-export tool to the TensorRT‑YOLO engine format.
2. Build the library (CMake) – enable BUILD_PYTHON=ON if you need the Python API.
3. Load the engine in code:

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

   or in C++:

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

4. Optional – call model.clone() (or det->clone()) to obtain a thread‑safe copy for multi‑threaded servers.
5. Profile – model.profile() returns throughput and latency metrics.

Why it matters – YOLO models are popular for real‑time vision, but getting the best speed on NVIDIA hardware usually requires hand‑crafted TensorRT plugins and careful CUDA tuning. TensorRT‑YOLO bundles those optimizations and exposes a clean API, letting developers focus on application logic rather than low‑level engine plumbing.

License & community – GPL‑3.0, with active issue tracking, a Bilibili tutorial series, and a sponsor page for donations.

---

Quick start (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

Now you can run the provided examples/detect script or write your own code as shown above.