【備忘録】PyTorch パフォーマンスチューニングガイド
以下のドキュメント、Blogを読んだので内容を備忘録的にメモしています。
1. num_workers>0
num_workers=0の場合、前のプロセスが完了しないと次のデータの読み込みが開始されないため、0以上の値にすることが推奨されている（増やすほどRAMの消費は大きくなる）。どの程度の数にするかは以下のLinkでは、4 * num_GPU, 8 or 16が推奨されている。
2. pin_memory=True
DataLoaderに固定のRAMを割り当て、そこからVRAMへデータを転送できるため、時間を節約できる。デフォルトはFalseなので、GPUを使う場合はTrueにすることを推奨。
3. ターゲットデバイスで直接Tensorを操作する
torch.Tensorが必要な時にはネイティブのPython, NumPyを使用して変換するのではなく、使用するデバイス上でtorch.Tensorを作成する（不要な変換、転送を避ける）。
# 使用するデバイスがGPUの場合
torch.rand(size, device=torch.device('cuda'))
torch.Tensor()の代わりにtorch.from_numpy(numpy_array)やtorch.as_tensor(others)を使う
torch.Tensor()は常にデータをコピーするため、ソースデバイスとターゲットデバイスがCPUの場合、torch.from_numpy(numpy_array)やtorch.as_tensor(others)を使えばデータのコピーを回避できる可能性がある。
データ転送とカーネル実行にオーバーラップがある場合はtensor.to(non_blocking=True)を使う
基本的に非同期転送によって実行時間を短縮できる。
for features, target in loader:
features = features.to('cuda:0', non_blocking=True)
target = target.to('cuda:0', non_blocking=True)
# ↓が同期ポイントで前の2行を待つ
output = model(features)
PyTorch JIT によってPointwise(Elementwise) の操作を単一のカーネルに融合する
PyTorch JITでは隣接するPointwiseの操作を単一のカーネルに融合して、メモリアクセス時間とカーネルの起動時間を償却できる（コンパイラーでまだ実装されていない融合の機会もある）。
Pointwiseの操作一覧
ネットワークのアーキテクチャデザインとバッチサイズを8の倍数に設定する
Nvidia GPUのTensorコアが行列の乗算で最高のパフォーマンスを出すため。特にNLPの場合、出力次元 (通常、vocabularyサイズ) などを確認すること。これはプロセスの種類 (forward passや勾配計算など) とcuBLAS のバージョンによっても異なる場合がある。
PyTorch AMPでFP16を使用する場合、FP16は8の倍数を必要とするため、通常は8の倍数が推奨される。
forward passにmixed precisionを使用してbackward passには使用しない
一部の操作はfloat32の精度を必要としないため、精度を低く設定するとメモリと実行時間の両方を節約できる。PyTorchのAMPは操作によって必要な精度（float16, float32）を自動的に適用する。
model.zero_grad()やoptimizer.zero_grad()の代わりにparameter.grad = Noneを使う
メモリ操作の回数を削減できる。
# PyTorch >= 1.7では以下のどちらかを使う
model.zero_grad(set_to_none=True)
optimizer.zero_grad(set_to_none=True)
推論 or 検証のための勾配計算をオフにする
モデルの出力のみを計算する場合は勾配計算は必要ないため、推論, 検証の時は勾配計算を無効にする。
# 以下のどちらかを使う
with torch.no_grad():
output = model(input)
@torch.no_grad()
def validation(model, input):
output = model(input)
return output
torch.backends.cudnn.benchmark = Trueにする
学習のループの前に実行すると自動チューナーがcuDNNの畳み込みを計算するためのアルゴリズムを最適化して高速化できる。入力サイズが頻繁に変わる場合は自動チューナーが頻繁にベンチマークする必要があり、パフォーマンスが低下することもある。
channels_lastメモリ形式を使用する
画像のような4D入力でNCHW形式ではなく、NHWC形式を使用する。チャンネルごとではなく、ピクセルごとに画像を保存していることになり、メモリ内のRGBレイヤーの各ピクセル値が近づく。
まだベータ段階で4DのNCHW形式のみがサポートされている。
Batch normalizationの直前にある畳み込み層のバイアスを無効にする
Batch normalizationの平均減算によって数学的にバイアスの影響が相殺されるため不要。モデル パラメーター、実行時間、およびメモリを節約できる。
DataParallelの代わりにDistributedDataParallelを使用する
マルチGPU環境では1ノードでもマルチプロセッシングが適用されるDistributedDataParallelの方がパフォーマンスが良い。
最後に
この記事ではまったく触れていないが、PyTorch2.0になってチューニングガイドがアップデートされていた。自分が2.0に移行したら内容を追っかけようと思います。
Discussion