【YOLOX】PyTorch版モデルをTensorRTへ変換する

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公式のスクリプトを使って簡単に変換できます。

1. 動作環境に入り、YOLOXリポジトリのルートをカレントディレクトリにします。

   cd /path/to/YOLOX
   

2. tools/trt.pyスクリプトを実行すると、PyTorch版モデルをTensorRT版に変換したモデルmodel_trt.pthが生成されます。

   python tools/trt.py -n yolox-m -c /path/to/yolox_m.pth
   

   上記コマンドの実行に成功すると、カレントディレクトリ(/path/to/YOLOX)直下にディレクトリYOLOX_outputs/yolox_mが生成され、そのディレクトリ下にmodel_trt.pthが出力されます。

   /path/to/YOLOX/YOLOX_outputs/yolox_m/
       ├ model_trt.engine
       └ model_trt.pth
   

動作確認

生成した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程度でしたので、期待通り高速化されたと言えそうです。

補足(環境構築)

YOLOX

git clone [e-mail]: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

TensorRT

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

https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/installing.html

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

torch2trt

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

その他のインストールオプションはhttps://github.com/NVIDIA-AI-IOT/torch2trt?tab=readme-ov-file#setupに記載されています。

参考

• README: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/demo/TensorRT/python/README.md#convert-model
• モデルの出力先に関する参照先:
  • https://github.com/Megvii-BaseDetection/YOLOX/blob/main/tools/trt.py#L70
  • https://github.com/Megvii-BaseDetection/YOLOX/blob/main/tools/trt.py#L46
  • https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/exp/build.py#L27
  • https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/exp/base_exp.py#L21