arutema47's blog

書いたり書かなかったり。

Pytorch高速化 (3) TensorRTで推論を10倍高速化

TLdr;

torch2trtというpytorchモデルをTensorRTに簡単に変換するライブラリを使い、Jetson nano+xavier上で画像認識とセグメンテーションの推論処理を10倍高速化できることを確認しました。

ただtorch2trtはカスタムモデルには対応していないため(resnetなどtorchvision標準モデルのみ)、自作モデルのTensorRT変換は大変だと思います。

他高速化シリーズ aru47.hatenablog.com

aru47.hatenablog.com

TensorRTとは

f:id:aru47:20201130141422p:plain https://amzn.to/3q9qrEK

https://amzn.to/37ndtL5

例えば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で走らせる事でエッジデバイスでも高速な推論が可能となる。

https://camo.qiitausercontent.com/d1d295922add81fb3bd160a2be4f091d2147b916/68747470733a2f2f71696974612d696d6167652d73746f72652e73332e61702d6e6f727468656173742d312e616d617a6f6e6177732e636f6d2f302f3137313931352f62653963613665332d383261612d663163372d303663362d3138616136633363616264612e706e67

引用:http://on-demand.gputechconf.com/gtcdc/2017/presentation/dc7172-shashank-prasanna-deep-learning-deployment-with-nvidia-tensorrt.pdf

その機能の一つとして強力なものにグラフ最適化というものがある。 図はInceptionモジュールのものであるが、通常のpytorch等はconvrelubatchnormと3ステップ別々の計算を行う。 ここで入出力のテンソルは毎回メモリから読み書きする時間も必要であり、計算時間は長くなってしまう。

そこでTensorRTのグラフ最適化(レイヤフュージョン)はDNNではCNN, relu, batchnormと3ステップ計算が頻出することを利用し、その計算をまとめた専用レイヤーを用意する。 そうすることでデータ読み書きと演算はあたかも1度しか行わずにconvrelubatchnormの計算を実行することができる。reluやbatchnormは単純な加算、クリッピングで表現できるため計算をまとめることは容易である(カーネルを用意さえあれば)。

f:id:aru47:20201128114742p:plain

引用:http://on-demand.gputechconf.com/gtcdc/2017/presentation/dc7172-shashank-prasanna-deep-learning-deployment-with-nvidia-tensorrt.pdf

例えばresnet50では等価的にレイヤ数を1/4に削減している。 この最適化は非常に強力であり、大幅な高速化が可能となる。

またPytorchでは対応していないINT8のTensorCoreを使った計算もTensorRTならば可能である。 FP16に対しメモリ量が半分になるため、データ伝送も2倍高速化することが期待できる(実際に大幅に高速化可能)

またTensorRTはV100などの一般的なGPUでも高速化が可能であり、10倍ほどの高速化が公式スライドで報告されている。

f:id:aru47:20201128114416p:plain

TesnorRTを気軽に試す

TensorRTを自分のネットワークに試すにはPytorchモデルをONNXフォーマットに変換→ONNXのモデルをTensorRTで読み込むというフローがあるが、ONNX変換はクセが強く、気軽に試すのは難しい。

そのためnvidiaの出しているtorch2trtという変換ツールを使うことでONNXを介さすに直接PytorchモデルをTensorRTに変換できます。 ただtorch2trtは対応しているレイヤは標準CNNのものがほとんどで(Resnetなど)、カスタム関数などが入っていると上手く変換はできませんでした。 例えば物体認識モデルの変換は難しいです。その場合は公開されているONNXモデルを使うのが良いです。

github.com

以下レポで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のページを読んでみて下さい。

github.com

画像認識

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.

画像セグメンテーション

f:id:aru47:20201128133902j:plain

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と大きく、量子化を積極的にやることで速度が大幅に上がることがわかりました。