TRTorchを使用してPyTorchの推論を高速化する

PyTorchの推論を高速化したいです。

TRTorchを使用すれば簡単にできるぞ

TRTorchについて

以降は下記のリンクを元に記述しています。

https://www.nvidia.com/en-us/on-demand/session/gtcfall20-a21864/

PyTorchはpython依存しているがTorchScriptを使用するとC++からも使用可能になります。モデル静的な構造に変換します。

PyTorchのモデルをTorch Script形式に変換してC++から呼ぶ一例は下記になります。

PyTorchのモデル、TorchScriptのモデルはTensorRTで推論高速化が可能です。

TensorRTの特徴は下記になります。

なるべく低い数値精度を保ったままモデルを変換します。

レイヤー、Tensor合成をしてGPUへのカーネル命令の呼び出しをまとめています。

ハードウェアごとに最適なカーネル命令が異なるのですが、自動的に最適なカーネル命令を取得します。

TRTorchのシステムアーキテクチャ

  • Lower:JITの内部表現をTensorRTで変換しやすいように変換する
  • 変換:主要なオペレーションと結果を確保
  • 変換:TensorRTに変換可能なオペレーションを変換
  • 変換:TensorRTで動作可能なオペレーションに変換されたグラフを返す
  • 実行:

まずグラフの値を固定します。

TensorRTで変換しやすい形に変更します。

TensorRTへの変換手順は下記です。

  • 変換コンテキストを作成
  • グラフの入力Tensorに基づいてTensorRT NetDefに入力を追加
  • ノードごとに入力を収集->主要な演算結果を入力して、主要な演算を評価
  • ノードと入力値を使用して、対応するノードをTensorRT NetDefに追加
  • 出力を追加

より詳細な内容は下記です。

https://nvidia.github.io/TRTorch/contributors/conversion.html#conversion

主要な演算の評価します。

主要操作をシミュレートした値を評価値マップに保存、Converterがをそれを使用してレイヤー設定を入力できるようにします。

下記の各ノードをTensorRTで対応するオペレーションに変換します。

実行

  • 変換が完了するとTensorRTエンジンが実行可能な新しいグラフに変換されます。
  • スタックが呼び出されると入力とエンジンをスタックからポップしてコンテキストマネージャーに渡し、エンジンを実行して出力をスタックにプッシュします。

下記の図の縦軸は速度向上率になっています。

パフォーマンスはJITでも高速になっていますが、TRTorchだとさらに高速化されています。

動作環境

下記の環境で動作確認しました。

動作環境は下記です。

ハードウェア

  • CPU: Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz
  • Memory: 64GB
  • GPU: Geforce 2080 Ti

ソフトウェア

  • OS: Ubuntu 18.04
  • cuda: バージョン11.0
  • Graphics Driver: バージョン450.51.06
  • docker: 19.03.12

TRTorchの導入

下記のコードをインストールします。

https://github.com/NVIDIA/TRTorch

TRTorchの直下に`WORKSPACE`というファイルがあるので、そこにあるcudnnとTensorRTをダウンロードします。

これらは2021年6月時点ではcuda11.1に依存しているので、cuda11.1のインストールが必要になります。

下記にcudaの一覧があるので、そこからダウンロードします。

https://developer.nvidia.com/cuda-toolkit-archive

https://github.com/NVIDIA/TRTorch

http_archive(
    name = "cudnn",
    build_file = "@//third_party/cudnn/archive:BUILD",
    sha256 = "98a8784e92862f20018d20c281b30d4a0cd951f93694f6433ccf4ae9c502ba6a",
    strip_prefix = "cuda",
    urls = [
        "https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.1.1.33/11.2_20210301/cudnn-11.2-linux-x64-v8.1.1.33.tgz",
    ],
)

http_archive(
    name = "tensorrt",
    build_file = "@//third_party/tensorrt/archive:BUILD",
    sha256 = "d3a1f478e304b48878604fac70ce7920fece71f9cac62f925c9c59c197f5d087",
    strip_prefix = "TensorRT-7.2.3.4",
    urls = [
        "https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/7.2.3/tars/TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz",

下記のディレクトリにcudnnとTensorRTのtarファイルをおきます。

third_party/dist_dir/x86_64-linux-gnu/

下記のようになります。

ls third_party/dist_dir/x86_64-linux-gnu/
cudnn-11.2-linux-x64-v8.1.1.33.tgz  TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz

bazelでビルドします。

bazel build //:libtrtorch --compilation_mode opt --distdir third_party/dist_dir/x86_64-linux-gnu/

Pythonで動作するようにインストールします。

cd py
pip install -r requirements.txt
python setup.py bdist_wheel
python3 setup.py install

必要なライブラリをインストールします。

pip install -U torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

jupyterで動作します。

jupyter-lab

物体検出モデルの動作確認

下記のコードで動作確認します。

https://github.com/NVIDIA/TRTorch/blob/master/notebooks/ssd-object-detection-demo.ipynb

ライブラリをインストールします。

%%capture 
%%bash
# Known working versions
pip install numpy==1.19 scipy==1.5.2 Pillow==6.2.0 scikit-image==0.17.2 matplotlib==3.3.0

必要なモデルのリストを取得します。

import torch

# List of available models in PyTorch Hub from Nvidia/DeepLearningExamples
torch.hub.list('NVIDIA/DeepLearningExamples:torchhub')

ssdモデルを取得します。

# load SSD model pretrained on COCO from Torch Hub
precision = 'fp32'
ssd300 = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd', model_math=precision);

推論用の画像データと後処理用のライブラリ`utils`を取得します。

# Sample images from the COCO validation set
uris = [
    'http://images.cocodataset.org/val2017/000000397133.jpg',
    'http://images.cocodataset.org/val2017/000000037777.jpg',
    'http://images.cocodataset.org/val2017/000000252219.jpg'
]

# For convenient and comprehensive formatting of input and output of the model, load a set of utility methods.
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd_processing_utils')

# Format images to comply with the network input
inputs = [utils.prepare_input(uri) for uri in uris]
tensor = utils.prepare_tensor(inputs, False)

# The model was trained on COCO dataset, which we need to access in order to
# translate class IDs into object names. 
classes_to_labels = utils.get_coco_object_dictionary()

モデルの推論と後処理をします。

0.4以上のスコアの検出結果を取得しています。

# Next, we run object detection
model = ssd300.eval().to("cuda")
detections_batch = model(tensor)

# By default, raw output from SSD network per input image contains 8732 boxes with 
# localization and class probability distribution. 
# Let’s filter this output to only get reasonable detections (confidence>40%) in a more comprehensive format.
results_per_input = utils.decode_results(detections_batch)
best_results_per_input = [utils.pick_best(results, 0.40) for results in results_per_input]

可視化するための関数を設定します。

from matplotlib import pyplot as plt
import matplotlib.patches as patches

# The utility plots the images and predicted bounding boxes (with confidence scores).
def plot_results(best_results):
    for image_idx in range(len(best_results)):
        fig, ax = plt.subplots(1)
        # Show original, denormalized image...
        image = inputs[image_idx] / 2 + 0.5
        ax.imshow(image)
        # ...with detections
        bboxes, classes, confidences = best_results[image_idx]
        for idx in range(len(bboxes)):
            left, bot, right, top = bboxes[idx]
            x, y, w, h = [val * 300 for val in [left, bot, right - left, top - bot]]
            rect = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='r', facecolor='none')
            ax.add_patch(rect)
            ax.text(x, y, "{} {:.0f}%".format(classes_to_labels[classes[idx] - 1], confidences[idx]*100), bbox=dict(facecolor='white', alpha=0.5))
    plt.show()

可視化して結果を確認します。

plot_results(best_results_per_input)

通常のPyTorchの推論速度

ベンチマークを測る関数を作成します。

Warm Upをして速度を測る準備をしています。

PyTorchで動作計測する場合はtorch.cuda.synchronize()をしないと正確に測れないので、こちらを使用します。

https://pytorch.org/docs/stable/notes/cuda.html#asynchronous-execution

import time
import numpy as np

import torch.backends.cudnn as cudnn
cudnn.benchmark = True

# Helper function to benchmark the model
def benchmark(model, input_shape=(1024, 1, 32, 32), dtype='fp32', nwarmup=50, nruns=1000):
    input_data = torch.randn(input_shape)
    input_data = input_data.to("cuda")
    if dtype=='fp16':
        input_data = input_data.half()
        
    print("Warm up ...")
    with torch.no_grad():
        for _ in range(nwarmup):
            features = model(input_data)
    torch.cuda.synchronize()
    print("Start timing ...")
    timings = []
    with torch.no_grad():
        for i in range(1, nruns+1):
            start_time = time.time()
            pred_loc, pred_label  = model(input_data)
            torch.cuda.synchronize()
            end_time = time.time()
            timings.append(end_time - start_time)
            if i%100==0:
                print('Iteration %d/%d, avg batch time %.2f ms'%(i, nruns, np.mean(timings)*1000))

    print("Input shape:", input_data.size())
    print("Output location prediction size:", pred_loc.size())
    print("Output label prediction size:", pred_label.size())
    print('Average batch time: %.2f ms'%(np.mean(timings)*1000))

通常のPyTorchの動作速度を計測します。

# Model benchmark without TRTorch/TensorRT
model = ssd300.eval().to("cuda")
benchmark(model, input_shape=(128, 3, 300, 300), nruns=1000)

下記の結果になりました。

Warm up ...
Start timing ...
Iteration 100/1000, avg batch time 526.72 ms
Iteration 200/1000, avg batch time 528.72 ms
Iteration 300/1000, avg batch time 531.57 ms
Iteration 400/1000, avg batch time 533.87 ms
Iteration 500/1000, avg batch time 535.29 ms
Iteration 600/1000, avg batch time 536.52 ms
Iteration 700/1000, avg batch time 537.99 ms
Iteration 800/1000, avg batch time 539.55 ms
Iteration 900/1000, avg batch time 540.22 ms
Iteration 1000/1000, avg batch time 541.96 ms
Input shape: torch.Size([128, 3, 300, 300])
Output location prediction size: torch.Size([128, 4, 8732])
Output label prediction size: torch.Size([128, 81, 8732])
Average batch time: 541.96 ms

TorchScriptモジュールを用いた推論速度

TorchScriptモジュールのモデルに変換します。

model = ssd300.eval().to("cuda")
traced_model = torch.jit.trace(model, [torch.randn((1,3,300,300)).to("cuda")])

変換したモデルを保存します。

# This is just an example, and not required for the purposes of this demo
torch.jit.save(traced_model, "ssd_300_traced.jit.pt")

ベンチマークを取ります。

# Obtain the average time taken by a batch of input with Torchscript compiled modules
benchmark(traced_model, input_shape=(128, 3, 300, 300), nruns=1000)

速度が若干、おそくなってしまっています。

Warm up ...
Start timing ...
Iteration 100/1000, avg batch time 537.86 ms
Iteration 200/1000, avg batch time 537.58 ms
Iteration 300/1000, avg batch time 536.58 ms
Iteration 400/1000, avg batch time 537.78 ms
Iteration 500/1000, avg batch time 539.76 ms
Iteration 600/1000, avg batch time 539.61 ms
Iteration 700/1000, avg batch time 539.65 ms
Iteration 800/1000, avg batch time 540.73 ms
Iteration 900/1000, avg batch time 541.80 ms
Iteration 1000/1000, avg batch time 542.49 ms
Input shape: torch.Size([128, 3, 300, 300])
Output location prediction size: torch.Size([128, 4, 8732])
Output label prediction size: torch.Size([128, 81, 8732])
Average batch time: 542.49 ms

TRTorchを用いた推論速度

TRTorchのモデルに変換します。FP16の数値精度で動作メモリサイズは2^20メモリ:1MBになります。

import trtorch

# The compiled module will have precision as specified by "op_precision".
# Here, it will have FP16 precision.
trt_model = trtorch.compile(traced_model, {
    "input_shapes": [(3, 3, 300, 300)],
    "op_precision": torch.half, # Run with FP16
    "workspace_size": 1 << 20
})

推論が正しく動作するか可視化をしてみます。

# using a TRTorch module is exactly the same as how we usually do inference in PyTorch i.e. model(inputs)
detections_batch = trt_model(tensor.to(torch.half)) # convert the input to half precision

# By default, raw output from SSD network per input image contains 8732 boxes with 
# localization and class probability distribution. 
# Let’s filter this output to only get reasonable detections (confidence>40%) in a more comprehensive format.
results_per_input = utils.decode_results(detections_batch)
best_results_per_input_trt = [utils.pick_best(results, 0.40) for results in results_per_input]

# Visualize results with TRTorch/TensorRT
plot_results(best_results_per_input_trt)

速度計測をします。バッチサイズはGPUメモリの都合上64に減らして動作確認しました。

batch_size = 32

# Recompiling with batch_size we use for evaluating performance
trt_model = trtorch.compile(traced_model, {
    "input_shapes": [(batch_size, 3, 300, 300)],
    "op_precision": torch.half, # Run with FP16
    "workspace_size": 1 << 20
})

benchmark(trt_model, input_shape=(batch_size, 3, 300, 300), nruns=1000, dtype="fp16")

結果は下記になります。他のケースでは128のバッチサイズのため、結果を2倍して考える必要がありますが。それでも6倍程度、高速になっています。

Warm up ...
Start timing ...
Iteration 100/1000, avg batch time 44.48 ms
Iteration 200/1000, avg batch time 44.45 ms
Iteration 300/1000, avg batch time 44.63 ms
Iteration 400/1000, avg batch time 44.83 ms
Iteration 500/1000, avg batch time 44.98 ms
Iteration 600/1000, avg batch time 44.93 ms
Iteration 700/1000, avg batch time 44.92 ms
Iteration 800/1000, avg batch time 44.97 ms
Iteration 900/1000, avg batch time 45.05 ms
Iteration 1000/1000, avg batch time 45.05 ms
Input shape: torch.Size([62, 3, 300, 300])
Output location prediction size: torch.Size([62, 4, 8732])
Output label prediction size: torch.Size([62, 81, 8732])
Average batch time: 45.05 ms

Close Bitnami banner
Bitnami