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 追蹤、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 腳本，或如上所示編寫您自己的程式碼。