概要
PytorchでInfenrece速度を上げるためのテクニック。NVIDIAのGPUを使う仮定です。思いつき次第、随時追加していきます。抜けているのがあればコメントして頂けると助かります。
メモリー系
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