TLdr;
torch2trtというpytorchモデルをTensorRTに簡単に変換するライブラリを使い、Jetson nano+xavier上で画像認識とセグメンテーションの推論処理を10倍高速化できることを確認しました。
ただtorch2trtはカスタムモデルには対応していないため(resnetなどtorchvision標準モデルのみ)、自作モデルのTensorRT変換は大変だと思います。
他高速化シリーズ aru47.hatenablog.com
TensorRTとは
例えばJetsonNanoなどは安価に入手できるエッジデバイスだが搭載しているGPUはローエンドなため大きい画像をリアルタイムで処理するのは難しい。
一般的なpytorch推論は
model.eval() with torch.no_grad(): for (data, target, _, _) in tqdm(loader): # data = data.half().cuda() # 半精度推論 logits = model(data)
と書けるが、これでも速度が不十分なことが多い。
TensorRTはnVidiaが提供している推論を高速化させるフレームワークである。モデルをTensorRTで走らせる事でエッジデバイスでも高速な推論が可能となる。
その機能の一つとして強力なものにグラフ最適化というものがある。
図はInceptionモジュールのものであるが、通常のpytorch等はconv
、relu
、batchnorm
と3ステップ別々の計算を行う。
ここで入出力のテンソルは毎回メモリから読み書きする時間も必要であり、計算時間は長くなってしまう。
そこでTensorRTのグラフ最適化(レイヤフュージョン)はDNNではCNN, relu, batchnormと3ステップ計算が頻出することを利用し、その計算をまとめた専用レイヤーを用意する。
そうすることでデータ読み書きと演算はあたかも1度しか行わずにconv
、relu
、batchnorm
の計算を実行することができる。reluやbatchnormは単純な加算、クリッピングで表現できるため計算をまとめることは容易である(カーネルを用意さえあれば)。
例えばresnet50では等価的にレイヤ数を1/4に削減している。 この最適化は非常に強力であり、大幅な高速化が可能となる。
またPytorchでは対応していないINT8のTensorCoreを使った計算もTensorRTならば可能である。 FP16に対しメモリ量が半分になるため、データ伝送も2倍高速化することが期待できる(実際に大幅に高速化可能)
またTensorRTはV100などの一般的なGPUでも高速化が可能であり、10倍ほどの高速化が公式スライドで報告されている。
TesnorRTを気軽に試す
TensorRTを自分のネットワークに試すにはPytorchモデルをONNXフォーマットに変換→ONNXのモデルをTensorRTで読み込むというフローがあるが、ONNX変換はクセが強く、気軽に試すのは難しい。
そのためnvidiaの出しているtorch2trtという変換ツールを使うことでONNXを介さすに直接PytorchモデルをTensorRTに変換できます。 ただtorch2trtは対応しているレイヤは標準CNNのものがほとんどで(Resnetなど)、カスタム関数などが入っていると上手く変換はできませんでした。 例えば物体認識モデルの変換は難しいです。その場合は公開されているONNXモデルを使うのが良いです。
以下レポでtensorRTをJetson上で試し、速度向上効果を見てみました。
TensorRT化はモデルと入力画像サイズを渡すことで簡単にできます。
# define input input_size = [1, 3, 256, 256] x = torch.zeros(input_size).cuda() # convert to tensorrt models model_trt = torch2trt(model, [x], fp16_mode=True, int8_mode=True, max_batch_size=1) # 精度によってモード切り替え
詳しい使い方はレポかtorch2trtのページを読んでみて下さい。
画像認識
resnet18 | resnet34 | resnet50 | |
---|---|---|---|
Raw | 11 | 12 | 16 |
FP32 | 3.8 | 5.6 | 9.9 |
FP16 | 2.1 | 3.3 | 4.4 |
INT8 | 1.7 | 2.7 | 3.0 |
Xavier上で数種類のネットワークの画像一枚に掛かる時間を調べました(ms)。rawはpytorchそのままで他はtensorRTで記載分解能で量子化したものです。NanoではInt8が何故か動きませんでした。 画像サイズは256x256.
画像セグメンテーション
fcn_resnet50 | fcn_resnet101 | deeplabv3_resnet50 | deeplabv3_resnet101 | |
---|---|---|---|---|
Raw | 200 | 344 | 281 | 426 |
FP32 | 173 | 290 | 252 | 366 |
FP16 | 36 | 57 | 130 | 151 |
INT8 | 21 | 32 | 97 | 108 |
こちらはセグメンテーションの結果です。画像サイズは512*512と大きく、量子化を積極的にやることで速度が大幅に上がることがわかりました。