「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」コンテナを起動します。



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

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

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

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

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

環境変数「TRT_DATA_DIR」のセット

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

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

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


ファイルの末尾に、以下を追加します。

「ctrl + o」を押下し、ファイルを上書き保存します。
「ctrl + x」を押下し、ファイルを閉じます。

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

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



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

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

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

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

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



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



コマンドオプション解説

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

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



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

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

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



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

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

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

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

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

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

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

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



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

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

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

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



コマンドの結果として、
「/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」と入力します。

参考ウェブサイト