Jetson NanoでTensorRTを使用したVGG16モデルによる画像判別

# はじめに

`TensorRT`の推論がスゴいという話なので勉強した。モデルはonnx-chainerを使って`chainer`から作成したONNX形式のVGG16モデルを用いる。`TensorRT`のサンプルが難しく理解するのに時間を要した。とにかくドキュメントとソースコード(C++, Python)を読みまくった結果「実はそんなに難しくないのでは・・・」と思い始めた。

本記事は「[Jetson Nanoでonnx-chainerを使ってONNX形式で出力](https://qiita.com/yamamo-to/items/55ebada2e02a199aae34)」の続編である。また`TensorRT`を使わない「[Jetson NanoでVGG16を試す](https://qiita.com/yamamo-to/items/d66845a1cd561bebaf86)」との比較でもある。

# TensorRTとは?

概要としては

- C++で作られたライブラリ
- 訓練済みのモデルを使用して推論を高速化
- APIとしてC++用とPython用を提供

といったところ。実行ステップは

1. モデルの読み込み
2. エンジンのビルド
3. 推論の実行

を行う。利用形態は

- C++もしくはPythonでAPIを直接叩く
- 好きなフレームワークでモデルを書いてONNX形式で出力し読み込ませる
- TensorRTをサポートするフレームワークを使って呼び出す

という感じである。今回は二つ目のONNX形式ということになる。

# 準備

`tensorrt`を`pipenv`の環境に入れるのが大変そうなので`pipenv`はここで一旦断念[^1]。ここからは`pip3`で頑張ることにする。意を決して`CuPy`も`Chainer`も`pip3`でインストールし直す。

また`PyCuda`が必要になるがインストールする前に環境変数`PATH`と`LD_LIBRARY_PATH`を設定する。

[^1]: 後で気が付いたが次のコマンドで行けた`export PYTHONPATH=/usr/lib/python3.6/dist-packages`

```bash:ターミナル
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`を作成した。

```python: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がでる。

```bash:ターミナル
python3 onnx2trt.py vgg16.onnx vgg16.trt
```

実行結果。

![onnx2trt.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/252224/83c2c207-d72c-61dc-aee7-646dc05b0435.png)

うまく実行できると`vgg16.trt`という新たに561MiBのエンジンをシリアライズしたファイルができる[^2]。

[^2]: 既にcaffemodel形式で528MiB、それをChainer用に変換したnpz形式が528MiB、さらにONNX形式に変換したので528MiBある。

# 推論の実行

色々なサンプルを参考に`trt_infer.py`を作成した。行っている事は
- 画像の読み込み
- エンジンの読み込み
- メモリの確保
- 画像のセット
- 推論
である。

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

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

```python: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()
```

このプログラムはコマンドライン引数にエンジンと画像を取るようにしているため、下記のように実行する。

```bash:ターミナル
python3 trt_infer.py vgg16.trt boss.jpg
```

実行結果。

![tensorrt_boss.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/252224/261e9e80-b5ae-8f28-37ac-9261f95523aa.png)

推論結果は752。ラベルはラケットらしい。

```
n04039381 racket, racquet
```
あれ？[前の結果](https://qiita.com/yamamo-to/items/d66845a1cd561bebaf86)は「障子」だったはずなんだが・・・。大丈夫なんだろうか。

失敗している感じがするので、適当に落としてきたリンゴの画像を試してみる。

![tensorrt_apple.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/252224/b2c93f27-eaec-cfc4-8612-fdf43418d006.png)

こちらは948と出た。推論結果のラベルはグラニースミス。

```
n07742313 Granny Smith
```

グラニースミスって何だろう、と調べたら「青リンゴ」のことだった。出来ているぽい?(これ赤リンゴのような・・・)

# 実行時間

最初の1回目の起動は1分以上掛かった。しかし2回目以降はとても速い。キャッシュでもされたんだろうか。`time`コマンドの結果は下記の通り[^3]。

[^3]: 経過時間のほとんどがエンジンのロード時間と思われることから、いずれ大量の画像で`TensorRT`を使用していない場合と使用した場合の実行時間比較を行いたい。

```bash:ターミナル
time python3 trt_infer.py vgg16.trt boss.jpg 
...
real    0m14.778s
user    0m6.460s
sys     0m3.464s
```

# さいごに

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

# 参考文献

- [NVIDIA Deep Learning SDK TensorRT Samples Support Guide](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html)
- [Migrating from TensorRT 4 to 5](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/migrationGuide.html)
- [Welcome to PyCUDA’s documentation!](https://documen.tician.de/pycuda/)
- [Fixstars Tech Blog:【TensorRTやってみた】](http://proc-cpuinfo.fixstars.com/2018/03/tensorrt-00/)
- [TensorRT 5.0使ってみた](https://qiita.com/dcm_sakai/items/0e13e2917adf55e92745)
- [Naren Dasan: TensorRT Python API](https://narendasan.com/posts/tensorrt-python-api.html)