laugh12321/TensorRT-YOLO

TensorRT‑YOLO – NVIDIA 하드웨어의 고성능 YOLO 추론

무엇인가 – TensorRT를 래핑하여 NVIDIA GPU 및 Jetson 장치에서 전체 YOLO 패밀리(v3‑v26, YOLO‑World, YOLO‑E, YOLO‑Master 등)를 실행하는 C++/Python 라이브러리입니다. 엔드투엔드 파이프라인(전처리 → 추론 → 후처리)을 최대한 빠르게 만들면서 API를 단순하게 유지하기 위해 사용자 정의 TensorRT 플러그인, CUDA 커널 및 CUDA‑Graph 지원이 함께 제공됩니다.

주요 기능

• 폭넓은 모델 지원 – 많은 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, 활발한 이슈 추적, 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 스크립트를 실행하거나 위와 같이 자신만의 코드를 작성할 수 있습니다.