「ONNX TensorRTを用いたYOLOv3による物体検出」のサンプルプログラムを試す

解説

NVIDIA社から「ONNX TensorRTを用いたYOLOv3による物体検出」のサンプルプログラムが提供されています。
試してみるサンプルプログラムは、「Object Detection With The ONNX TensorRT Backend In Python」です。

試してみるサンプルプログラムは、Jetson Orin Nano の「/usr/src/tensorrt/samples/python/yolov3_onnx」にあります。
（GitHubでも公開されています。「TensorRT/samples/python/yolov3_onnx at main · NVIDIA/TensorRT · GitHub」）

今回、サンプルプログラムを、Jetsonに最適化された機械学習開発用のコンテナ「L4T ML」上で試してみることとします。

機械学習開発用のコンテナの起動

「Dockerコンテナを実行してJupyterLabにアクセス」に従い、コンテナ起動スクリプトファイルを作成します。

以下のコマンドを実行し、「L4T ML」コンテナを起動します。

./docker_ml_run.sh

以降の操作は、起動したコンテナ上にいるTerminalで行ないます。

サンプルプログラムのユーザーのホームディレクトリ下へのコピー

サンプルプログラムは、「/usr/src/tensorrt/samples/python/yolov3_onnx」にありますが、このディレクトリはユーザーの書き込み権限がないため、プログラムの実行はできません。ユーザーのホームディレクトリ下へコピーします。

別ディレクトリにあるプログラムモジュールも使用するので、「/usr/src/tensorrt/samples/python/」以下をユーザーのホームディレクトリ下にコピーします。

以下のコマンドを実行し、コピー先のディレクトリを作成し、サンプルプログラムをコピーします。（コンテナ上の「/work」ディレクトリは、ホスト上の「~/work」ディレクトリをマウントしたディレクトリです。）

mkdir -p /work/tensorrt/samples cp -r /usr/src/tensorrt/samples/python/ /work/tensorrt/samples/python/

環境変数「TRT_DATA_DIR」のセット

データファイルを保存するときのベースディレクトリを環境変数で指定します。（プログラムの引数として指定することもできますが、環境変数で指定する方が一度の指定で済むので楽です）

以下のコマンドを実行し、環境変数「TRT_DATA_DIR」をセットします。環境変数の値は、「/work/tensorrt/」とします。

export TRT_DATA_DIR=/work/tensorrt/

参考）exportコマンドは、環境変数をセットします。再起動や再ログインすると、リセットされます。恒久的にセットする場合は、「~/.profile」に記述します。（コンテナは、「--rm」のないコンテナ起動コマンドで起動している必要があります）: 以下のコマンドを実行し、「~/.profile」を開きます。
nano ~/.profile

ファイルの末尾に、以下を追加します。
export TRT_DATA_DIR=/work/tensorrt/
「ctrl + o」を押下し、ファイルを上書き保存します。
「ctrl + x」を押下し、ファイルを閉じます。

以下のコマンドを実行し、「.profile」の変更を反映させます。
source ~/.profile

以下のコマンドを実行し、環境変数を確認します。

printenv

設定されている環境変数が表示されます。「TRT_DATA_DIR=/work/tensorrt/」がセットされました。

必要ファイルのダウンロード

試してみるサンプルプログラムが必要とするファイルは、
「/work/tensorrt/samples/python/yolov3_onnx/download.yml」
で指定されています。
「download.yml」ファイルを確認すると「yolov3.cfg」「yolov3.weights」「dog.jpg」の３つのファイルが指定されています。

「download.yml」ファイルで指定されているファイルのダウンロードは、
「/work/tensorrt/samples/python/downloader.py」
で行ないます。

以下のコマンドを実行し、「downloader.py」のディレクトリに移動します。

cd /work/tensorrt/samples/python/

以下のコマンドを実行し、「downloader.py」に必要なPythonパッケージをインストールします。

pip3 install --upgrade -r requirements.txt

コマンドオプション解説

--upgrade: 既にインストールされているPythonパッケージについて、アップグレードできるPythonパッケージについてはアップグレードする。（PyYAMLが既にインストールされていましたが、バージョンが3.12と古いバージョンでした。このオプションを指定しない場合、downloader.pyの実行において、エラーを生じました。）
-r: requirements.txt で指定されたPythonパッケージをインストールします。

以下のコマンドを実行し、試してみるサンプルプログラムが必要とするファイルをダウンロードします。

python3 downloader.py -f /work/tensorrt/samples/python/yolov3_onnx/download.yml

コマンドの結果として、
「/work/tensorrt/samples/python/yolov3_onnx/」直下に、「yolov3.cfg」「yolov3.weights」「dog.jpg」の３つのファイルがダウンロードされます。

サンプルプログラムに必要なPythonパッケージのインストール

以下のコマンドを実行し、サンプルプログラムのディレクトリに移動します。

cd /work/tensorrt/samples/python/yolov3_onnx

以下のコマンドを実行し、サンプルプログラムに必要なPythonパッケージをインストールします。

pip3 install --upgrade -r requirements.txt

サンプルプログラムの実行

サンプルプログラムを実行していきます。

サンプルプログラムは以下の段取りで、物体検出を行ないます。

YOLOv3学習済みモデルファイルをONNX形式ファイルに変換します。
ONNX形式ファイルからTensorRTエンジンをビルドします。
TensorRTを用いて物体検出の推論を実施します。

YOLOv3学習済みモデルファイルのONNX形式ファイルへの変換

YOLOv3学習済みモデルファイル「yolov3.weights」は、darknet形式で、TensorRTは読み込むことができないため、TensorRTが読み込むことができる、ONNX形式のファイルに変換します。
ONNX : Open Neural Network eXchange。

以下のコマンドを実行し、YOLOv3学習済みモデルファイルをONNX形式ファイルに変換します。

python3 yolov3_to_onnx.py

コマンドの結果として、
「/work/tensorrt/samples/python/yolov3_onnx/」直下に、「yolov3.onnx」ファイルが出来ます。

TensorRTエンジンのビルドと推論の実施

ONNX形式ファイルからTensorRTエンジンをビルドし、TensorRTを用いて物体検出の推論を実施します。

以下のコマンドを実行し、TensorRTエンジンのビルドと、物体検出の推論を実施します。

python3 onnx_to_tensorrt.py

コマンドの結果として、
「/work/tensorrt/samples/python/yolov3_onnx/」直下に、「dog_bboxes.png」ファイルが出来ます。

物体検出結果（dog_bboxes.png）

「dog」「bicycle」「truck」が検出されました。

TensorRTエンジンのビルドには時間を要しますが、TensorRTエンジンは再利用可能です。
上記コマンドをはじめて実行したときは、TensorRTエンジンのビルド（とTensorRTエンジンファイル出力）が行なわれ、処理完了に、4分ほど時間を要しました。
上記コマンドを再度実行したときは、TensorRTエンジンのビルドは行なわれず（代わりにTensorRTエンジンファイルの読み込みが行なわれます）、処理完了に、3秒ほどしか時間を要しませんでした。

コンテナの終了

コンテナを終了するには、起動したコンテナ上にいるTerminalで「exit」と入力します。