概要

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が連続して配置されている方がメモリアクセス効率が高い。

概要

メモリー系

量子化(Quantization)

モデルの設定と設計

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

Inference時の設定

cuDNN

メモリーレイアウト

Reference