はじめに
TensorRTの推論がスゴいという話なので勉強した。モデルはonnx-chainerを使ってchainerから作成したONNX形式のVGG16モデルを用いる。TensorRTのサンプルが難しく理解するのに時間を要した。とにかくドキュメントとソースコード(C++, Python)を読みまくった結果「実はそんなに難しくないのでは・・・」と思い始めた。
本記事は「Jetson Nanoでonnx-chainerを使ってONNX形式で出力」の続編である。またTensorRTを使わない「Jetson NanoでVGG16を試す」との比較でもある。
TensorRTとは?
概要としては
- C++で作られたライブラリ
- 訓練済みのモデルを使用して推論を高速化
- APIとしてC++用とPython用を提供
といったところ。実行ステップは
- モデルの読み込み
- エンジンのビルド
- 推論の実行
を行う。利用形態は
- C++もしくはPythonでAPIを直接叩く
- 好きなフレームワークでモデルを書いてONNX形式で出力し読み込ませる
- TensorRTをサポートするフレームワークを使って呼び出す
という感じである。今回は二つ目のONNX形式ということになる。
準備
tensorrtをpipenvの環境に入れるのが大変そうなのでpipenvはここで一旦断念1。ここからはpip3で頑張ることにする。意を決してCuPyもChainerもpip3でインストールし直す。
またPyCudaが必要になるがインストールする前に環境変数PATHとLD_LIBRARY_PATHを設定する。
export PATH=/usr/local/cuda-10.0/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-10.0/lib64:$LD_LIBRARY_PATH
pip3 install pycuda --user
pip3 install cupy --user
pip3 install chainer --user
pip3 install chainercv --user
pip3 install pillow --user
ONNX形式からエンジンのビルド
ONNX形式を読み込んでエンジンをビルドする。ビルドしたエンジンはシリアライズして保存する。これを行うonnx2trt.pyを作成した。
import argparse
import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
def build_engine(onnx_file_path, engine_file_path):
with trt.Builder(TRT_LOGGER) as builder:
builder.max_workspace_size = 1 << 20
builder.max_batch_size = 1
with builder.create_network() as network:
with trt.OnnxParser(network, TRT_LOGGER) as parser:
with open(onnx_file_path, 'rb') as model:
parser.parse(model.read())
engine = builder.build_cuda_engine(network)
with open(engine_file_path, 'wb') as f:
f.write(engine.serialize())
def main():
parser = argparse.ArgumentParser()
parser.add_argument('onnx_file_path')
parser.add_argument('engine_file_path')
args = parser.parse_args()
build_engine(args.onnx_file_path, args.engine_file_path)
if __name__ == '__main__':
main()
このファイルの使い方は下記の通り。なおスワップを有効にしておかないとmemory errorがでる。
python3 onnx2trt.py vgg16.onnx vgg16.trt
実行結果。
うまく実行できるとvgg16.trtという新たに561MiBのエンジンをシリアライズしたファイルができる2。
推論の実行
色々なサンプルを参考にtrt_infer.pyを作成した。行っている事は
- 画像の読み込み
- エンジンの読み込み
- メモリの確保
- 画像のセット
- 推論
である。
当たり前の話だが、TensorRTを使わないモデルと入力・出力を合わせなければならない。
入力に関して言えば、前回Chainerを使って画像をVGG16の形式に変換したので、今回もChainerを用いて入力とした。出力については、1000個のラベルの中からSoftmaxを通した後の値が配列に格納される。
import argparse
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from PIL import Image
from chainer import links as L
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
def load_engine(filename):
with open(filename, 'rb') as f:
with trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
def load_image(filename):
img = Image.open(filename)
img = L.model.vision.vgg.prepare(img)
img = img[np.newaxis]
img = np.array(img, dtype=np.float32)
return img
def main():
parser = argparse.ArgumentParser()
parser.add_argument('engine_file_path')
parser.add_argument('image_file_path')
args = parser.parse_args()
engine = load_engine(args.engine_file_path)
img = load_image(args.image_file_path)
with engine.create_execution_context() as context:
# allocate memory
input_shape = engine.get_binding_shape(0)[:]
output_shape = engine.get_binding_shape(1)[:]
h_input = cuda.pagelocked_empty(shape=input_shape, dtype=np.float32)
h_output = cuda.pagelocked_empty(shape=output_shape, dtype=np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
bindings = [int(d_input), int(d_output)]
# set input image
h_input[:] = img
# do inference
stream = cuda.Stream()
cuda.memcpy_htod_async(d_input, h_input, stream)
context.execute_async(bindings=bindings, stream_handle=stream.handle)
cuda.memcpy_dtoh_async(h_output, d_output, stream)
stream.synchronize()
# show result
print(h_output)
print(np.argmax(h_output))
if __name__ == '__main__':
main()
このプログラムはコマンドライン引数にエンジンと画像を取るようにしているため、下記のように実行する。
python3 trt_infer.py vgg16.trt boss.jpg
実行結果。
推論結果は752。ラベルはラケットらしい。
n04039381 racket, racquet
あれ？前の結果は「障子」だったはずなんだが・・・。大丈夫なんだろうか。
失敗している感じがするので、適当に落としてきたリンゴの画像を試してみる。
こちらは948と出た。推論結果のラベルはグラニースミス。
n07742313 Granny Smith
グラニースミスって何だろう、と調べたら「青リンゴ」のことだった。出来ているぽい?(これ赤リンゴのような・・・)
実行時間
最初の1回目の起動は1分以上掛かった。しかし2回目以降はとても速い。キャッシュでもされたんだろうか。timeコマンドの結果は下記の通り3。
time python3 trt_infer.py vgg16.trt boss.jpg
...
real 0m14.778s
user 0m6.460s
sys 0m3.464s
さいごに
TensorRTは実はそんなに難しくないことが分かった。サンプルが難しいのと、ドキュメントが最新のバージョンについていけていないのが気になった(TensorRTとPyCudaの両方)。またTensorRTを使わない場合と使った場合とで推論結果が異なることもあるというのがわかった。そろそろ上位5つを表示した方がいいだろうか。