TensorRT5がリリースされたので試してみた
タイトルの通り正式版としてTensorRT 5.0.2.6が出たということでPython APIを試してみた。
やったことはTensorFlowのモデルを変換して動かしてみるところまで。
公式ドキュメントを参考に進めたけど情報が盛りだくさんでちょっとわかりにくかったので改めてまとめてみた。
コードはこちら。
TensorRTのPythonインタフェースに関する公式ドキュメントはこちら。
環境
動作させた環境は以下で、一通りインストールはできている前提。
OS: Ubuntu16.04
Python: 3.5
GPU: GTX 1080 Ti
CUDA: 10.0
cuDNN: 7.3.0
TensorRT インストール
公式ドキュメントを参考に以下からダウンロード。 docs.nvidia.com
Ubuntuに関しては14.04から18.04まで、14.04と16.04はCUDA9.0と10.0をそれぞれサポート。18.04はCUDA10.0だけ。
システムにインストールしてもよければdebパッケージのほう、自分で管理したければtarのほうを選ぶ。
自分はPython関連はvirtualenvで管理したかったのでtarのほうのwhlをpipで、それ以外はdebパッケージから入れた。
Pythonパッケージインストール
今回のコードを動かすのに必要なPythonパッケージのインストール
$pip install Pillow pycuda numpy
モデルの準備
MobilenetV1の学習済みモデルをTensorFlow modelsからダウンロード。
モデルの変換
Pythonパッケージがインストールできていればconvert-to-uff
というコマンドが使えるようになっているので、以下の通りuff形式に変換する。
$convert-to-uff mobilenet_v1_1.0_224_frozen.pb -o mobilenet_v1_1.0_224.uff
実行
$python run_tensorrt.py cat.jpg build engine... allocate buffers load input Inference time 0: 0.8375644683837891 [msec] Inference time 1: 0.7922649383544922 [msec] Inference time 2: 0.8244514465332031 [msec] Inference time 3: 0.7870197296142578 [msec] Inference time 4: 0.8032321929931641 [msec] Inference time 5: 0.7877349853515625 [msec] Inference time 6: 0.8432865142822266 [msec] Inference time 7: 0.8103847503662109 [msec] Inference time 8: 0.8285045623779297 [msec] Inference time 9: 0.8194446563720703 [msec] Classification Result: 1 tiger cat 0.45609813928604126 2 cougar 0.2968028783798218 3 Persian cat 0.24455446004867554 4 leopard 0.00040596097824163735 5 cheetah 0.00032873026793822646
変換したuffモデルを読み込んで引数で与えた画像に対する推論を10回実行している。 だいたい 0.8 msec / image = 1250 fps。 画像のロードとかGPUへのデータ転送は測定に入れてないにしても速い。
コードは以下。
基本的にはサンプルコードそのままで、時間測定のために非同期での実行をやめたくらい。
あとがき
DTYPE
のところをfloat16
にして試してみたけど、1080Tiは非対応ということで警告が出て速度も速くはならなかった。
int8
はこのモデルのままではダメで、int8
に量子化したTensofFlowのモデルをuffにすればできるかもしれない。この辺りはまた後でやってみたい。
あと、認識結果がTensorFlowと比べて悪くなってる気がするので、まだどこかおかしいかもしれない。
TensorFlowがchannel lastだけどTensorRTがchannel firstなところが気になって、register_input
のorderを変えたけど特に変化はなく、入力のtransposeを外すとさらに変な結果になってしまった。
完全におかしいわけじゃないので、前処理が違うとかかもしれない。
知っている人がいたら教えてほしい。