はじめに

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はここで一旦断念¹。ここからは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を作成した。

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のエンジンをシリアライズしたファイルができる²。

推論の実行

色々なサンプルを参考にtrt_infer.pyを作成した。行っている事は

画像の読み込み
エンジンの読み込み
メモリの確保
画像のセット
推論
である。

当たり前の話だが、TensorRTを使わないモデルと入力・出力を合わせなければならない。

入力に関して言えば、前回Chainerを使って画像をVGG16の形式に変換したので、今回もChainerを用いて入力とした。出力については、1000個のラベルの中からSoftmaxを通した後の値が配列に格納される。

trt_infer.py

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コマンドの結果は下記の通り³。

ターミナル

time python3 trt_infer.py vgg16.trt boss.jpg 
...
real    0m14.778s
user    0m6.460s
sys     0m3.464s

さいごに

TensorRTは実はそんなに難しくないことが分かった。サンプルが難しいのと、ドキュメントが最新のバージョンについていけていないのが気になった(TensorRTとPyCudaの両方)。またTensorRTを使わない場合と使った場合とで推論結果が異なることもあるというのがわかった。そろそろ上位5つを表示した方がいいだろうか。

参考文献

後で気が付いたが次のコマンドで行けたexport PYTHONPATH=/usr/lib/python3.6/dist-packages ↩
既にcaffemodel形式で528MiB、それをChainer用に変換したnpz形式が528MiB、さらにONNX形式に変換したので528MiBある。 ↩
経過時間のほとんどがエンジンのロード時間と思われることから、いずれ大量の画像でTensorRTを使用していない場合と使用した場合の実行時間比較を行いたい。 ↩