TF-TRT

NVIDIA Jetson Nano 開発者キットで TensorRT サンプル・プログラムを試す で TensorRT を Jetson Nano で利用する方法について紹介させていただきましたが、TensorRT を利用する方法はもう一つ存在します。TensorRT は TensorFlow に統合されていますので TensorFlow から TensorRT の機能を利用することができます。この仕組みは TF-TRT と呼ばれ、TensorRT 対応のニューラルネットワーク層は TensorRT で最適化され、そうでない層は TensorFlow でそのまま処理されます。ニューラルネットワーク・モデルを TensorRT 向けにすべて変換する方法に比べると多少パフォーマンスは劣りますが簡単に利用できるのが利点です。

TensorFlow/TensorRT Models on Jetson

TF-TRT を利用した物体認識と物体検出のサンプルが NVIDIA-AI-IOT/tf_trt_models で公開され、それぞれ Jupyter ノートブック上で試すことができます。ここでは物体検出の Jupyter ノートブック tf_trt_models/examples/detection/detection.ipynb を試してみましょう。

まず、Setup で説明されているとおりに必要なモジュールをインストールします。Python3 で試す場合、ステップ２にインストールしているモジュール python-pip と python-matplotlib はそれぞれ python3-pip と python3-matplotlib に変更が必要です。

Setup で説明されているモジュールに加え、Jupyter のインストールも必要です。

$ sudo pip3 install jupyter

また、TensorFlow で学習済みのモデルに対して TensorRT による最適化をかける過程で多くのメモリを必要とします。以下のとおりスワップ領域を割り当ててこの問題を乗り切ります。¹

$ sudo fallocate -l 4G /var/swapfile
$ sudo chmod 600 /var/swapfile
$ sudo mkswap /var/swapfile
$ sudo swapon /var/swapfile
$ sudo bash -c 'echo "/var/swapfile swap swap defaults 0 0" >> /etc/fstab'

インストールが完了したら Jupyter Notebook を起動します。

$ jupyter notebook

ウェブブラウザが自動的に起動します。

examples -> detection と移動して detection.ipynb をクリックします。
後は Jupyter ノートブック内の Python スクリプトを順に実行することで TF-TRT の仕組みを学ぶことができます。ここで行われる処理は以下のとおりです。

1. TensorFlow で学習済みのモデルからフローズン・グラフ ² を生成
2. フローズン・グラフを TensorRT で最適化
3. 最適化された TensorFlow グラフを TensorFlow にロード
4. 推論（物体検出）を実行

カメラ画像から物体検出

上記 detection.ipynb Jupyter ノートブックの実行中に TensorRT で最適化された TensorFlow グラフ（モデル）が tf_trt_models/examples/detection/data/ssd_inception_v2_coco_trt.pb ファイルとして保存されますが、このモデルを使ってカメラ画像のほぼリアルタイム物体検出を試してみました。

コードを tsutof/tf-trt-ssd で公開中です。

モデル読み込みに時間がかかるため起動に３分程待つ必要があります。

以上です。