Melody-Zhou/tensorRT_Pro-Vision

This repository is forked from shouxieai/tensorRT_Pro, extending it to support a wide range of vision models with high-performance TensorRT C++ inference.

解決する課題

NVIDIA GPU上で多数のビジョンモデルをデプロイする場合、通常はタスクごとにプロジェクトを維持する必要があり、CPUとGPU間のデータ転送がボトルネックとなります。本フレームワークは、20種類以上のタスク(検出、分類、セグメンテーション、姿勢推定、回転ボックス、深度推定、OCR、車線検出、マルチオブジェクトトラッキング)を単一のC++実行ファイルに統合し、すべての前処理と後処理をGPU上で実行することでこの問題を解決します。

仕組み

本プロジェクトはTensorRT C++ APIに基づいて構築されています。すべてのモデルは非同期推論スケルトンを共有しており、違いは前処理と後処理のみです。モデルをONNX(動的バッチ対応)としてエクスポートし、ワークスペースフォルダに配置して、ワークスペースから ./pro を実行するだけです。フレームワークがTensorRTエンジン(FP32、FP16、またはINT8)をコンパイルして実行します。前処理(リサイズ、正規化)と後処理(デコード、NMS)はCUDAカーネルとして実装されているため、データはGPU上に留まります。また、必要なモデル向けにカスタムTensorRTプラグイン(LayerNorm、DCNv2など)も含まれています。新しいタスクを追加するには、Inferインターフェース、preprocess()関数、worker()関数、および後処理用CUDAカーネルの3つを実装し、main.cppでタスクを登録するだけです。

対象ユーザー

NVIDIA GPU上で低遅延かつ高スループットなビジョンモデルのデプロイが必要なC++開発者やエンジニア、特に複数の異なるビジョンタスクを含む製品を開発している方に最適です。CUDA、TensorRT、およびMakefileやCMakeのパス変更に慣れている必要があります。また、YOLOのバリエーション、OCR、車線検出、深度推定などを単一のコードベースから迅速にベンチマークしたり出荷したりしたい研究者にも役立ちます。

ハイライト

  • 20種類以上のモデルを単一の実行ファイルで:タスク名で切り替え可能、タスクごとのプロジェクト管理が不要。
  • フルGPUパイプライン:リサイズ、正規化、デコード、NMS、深度リサイズはすべてCUDAカーネルで実行され、CPU↔GPU間の往復を回避。
  • TensorRT 8.x/10.x 適応型:条件付きコンパイルにより、protobufなしで両方のAPI世代をサポート。
  • 動的バッチと3つの精度:FP32、FP16、INT8をサポートし、EntropyおよびMinMaxキャリブレータを内蔵。
  • モジュール式拡張:タスクの追加は preprocess()、worker()、および後処理カーネルの実装のみで完了。
  • 幅広いタスク対応:検出(YOLOv3–v13, v26, DETR)、分類、セグメンテーション、姿勢推定、OBB、深度推定、OCR、車線検出、ByteTrackトラッキング。

関連

  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト