YOLOXの処理を高速化する目的で、既存のPyTorch版モデルをTensorRT版に変換してみました。

前提

YOLOXリポジトリをcloneし、TensorRTとtorch2trtをインストールした動作環境があること。 (環境構築については、補足を参照してください。)
YOLOXのPyTorch版のモデル(重みファイル)が手元にあること。
- 重みファイルは下記README中の一覧表 (weights) からダウンロードできます。
  https://github.com/Megvii-BaseDetection/YOLOX?tab=readme-ov-file#benchmark
- 以降では、yolox_m.pthをダウンロードした前提で実行コマンドを記載します。

PyTorch→TensorRT変換手順

動作環境が用意できていれば、YOLOX公式のスクリプトを使って簡単に変換できます。

生成したTensorRT版モデルを使って、YOLOXのデモ用スクリプトを実行します。

python tools/demo.py image -n yolox-m --trt --save_result

デモに成功すると、推論結果の可視化画像がvis_res/YYYY_mm_dd_HH_MM_SS/dog.jpgとして出力されます。

/path/to/YOLOX/YOLOX_outputs/yolox_m/
    ├ model_trt.engine
    ├ model_trt.pth
    └ vis_res/
        └ YYYY_mm_dd_HH_MM_SS/
            └ dog.jpg

推論処理にかかった処理時間は下記のように出力されます。筆者の環境では、約12FPSの処理速度でした。

Infer time: 0.0846s

同様にしてPyTorch版でデモを試した際は8FPS程度でしたので、期待通り高速化されたと言えそうです。

git clone git@github.com:Megvii-BaseDetection/YOLOX.git
cd YOLOX
pip install -v -e .

PyTorch版モデルでデモを試したい場合は、下記のDemoが参考になります。

https://github.com/Megvii-BaseDetection/YOLOX/tree/main?tab=readme-ov-file#quick-start

下記リンクの手順に従ってインストールできます。

debファイルからインストールする等いくつかのインストールオプションがありますが、いずれの方法でもNVIDIAアカウントが必要です。

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install