Pytorchで実行速度(inference speed)を上げる方法(GPU)

概要

PytorchでInfenrece速度を上げるためのテクニック。NVIDIAのGPUを使う仮定です。思いつき次第、随時追加していきます。抜けているのがあればコメントして頂けると助かります。

メモリー系

• GPUのSynchronizationを避ける

print(cuda_tensor)
cuda_tensor.item()
memory copies: tensor.cuda(), cuda_tensor.cpu() and equivalent tensor.to(device) calls
cuda_tensor.nonzero()
python control flow which depends on results of operations performed on cuda tensors e.g. if (cuda_tensor != 0).all()

• 非同期のデータ転送 -> data.to(device,non_blocking=True)
• TensorをGPU上に直接作る -> torch.rand(size,device="cuda")

量子化(Quantization)

• FP16を使う -> ほぼ全てのモデルが、ほとんど精度を落とさず高速化できる
• INT8を使う -> GPUでのINT8はTensorRTのみ。

モデルの設定と設計

• NormalizationはBatchNormを使う -> Inference時に平均と分散を計算しなくて良い。
• Normalizationの前のConvolutionのbiasをFalseにする -> Sequential(Conv2d(bias=False),BatchNorm2d())
• Layer Fusionをする -> Conv BatchNormは実行時に一つのカーネルに出来る。
• Reparametrizationをする -> 並列するConvolution層を実行時に一つのカーネルに出来る。
• channelのサイズは8の倍数(FP32,FP16)または32の倍数(INT8)にする -> Nvidiaのデータformatが8または32channelを1ブロックとしている為。端数が出ると空いたメモリをpaddingするらしい。
• MAC(Memory Access Cost)を減らす -> レイヤーやChannelが増えるとメモリーにアクセスする時間が割とかかるらしい。
• separable convolutionを使わない -> separable convolutionはパラメーターが少なくなるが、計算スピードは早くならないことが多い(レイヤーが増えるとMACが上がる+channelサイズが8の倍数ではない)。

変換またはコンパイルする

• TorchScriptに変換する(torch.jit.script) -> C++のようにコードをコンパイルする事で高速化。Pytorch1.12からNVFuserによりDynamic Shapeがサポートされた。
• モデルをFreezeさせる(torch.jit.freeze) -> 学習時に必要な機能を無くすことで、高速化。
• torch.compile -> コードを変えずにモデルを簡単にcompile出来る。
• TensorRTを使う -> Torch-TensorRTというPytorchのラッパーを使うと簡単。

Inference時の設定

• eval modeにする -> model.eval()
• gradientの計算をoffにする -> with toch.no_grad()を呼ぶ
• torch.inference_mode()を使う -> torch.no_grad()はgradientを計算しないのに対して、gradientをset出来なくなる。余計なメモリやoverheadを削減出来る。

cuDNN

• cudnn Auto-Tunerを使う -> torch.backends.cudnn.benchmark = True
• cudaのversionを最新にする。versionによってperformanceが結構違ったりする。

メモリーレイアウト

• torch.channels_lastを使う -> channel方向に畳込みをするので、channelが連続して配置されている方がメモリアクセス効率が高い。

Reference