本記事では画像認識系の国際会議 ECCV2020のチュートリアルでNVIDIAが発表した資料 PYTORCH PERFORMANCE TUNING GUIDE の内容をまとめるとともに、理解の助けになるような関連情報を参照します。 PyTorchは簡単にニューラルネットワークの実装のを容易さだけでなく、処理速度にも注意を払って開発が進められています。 プログラムにほんの少し修正を加えるだけで高速化できる可能性があります。 本記事を読み、実践することで、手元のPyTorchプログラム (特にGPUを使った学習処理) を高速化できる可能性があります。

目次

概要

チュートリアルで紹介されている内容は以下のような特徴があります。

  • 単純な方法で学習を高速化できる
  • 数行の変更で済む
  • 自動混合精度学習 (Automatic Mixed Precision; AMP) と併用できる

以降は、PyTorchのバージョンは1.6.0を想定します。

データ読み込みを非同期＆並列化する

データの読み込みが処理速度のボトルネックになっている場合、データ読み込みの並列化によって処理の高速化が期待されます。 PyTorchではDataLoaderクラスがデータ読み込みに関する処理を提供しています。 DataLoaderは引数のnum_workersに0以上の値を設定することで、非同期・並列化に対応できます。

また、非同期・並列化とは直接関係ありませんが、GPUを使った学習の場合、 DataLoaderの引数にpin_memory=Trueを与えると、CPUからGPUへのデータの転送が高速化される可能性があります1。

cuDNNのautotunerを利用する

cuDNNは畳み込み層に関して複数のアルゴリズムを実装しています。 torch.backends.cudnn.benchmark = True とすることで、PyTorchが自動で、処理速度の観点でハードウェアに適したアルゴリズムを選択してくれます2。

バッチサイズを増やす

GPUのメモリに余裕がある限りはバッチサイズを大きくした方が学習の処理が高速化できます。混合精度学習を適用すれば、さらにメモリに余裕ができるので、バッチサイズをより大きくできます。

また、バッチサイズを大きくしたら、学習率、warmupのステップ数、weight decayなどの調整を忘れないようにしましょう。もしくはoptimizerを大きなバッチサイズでの学習を想定したアルゴリズムであるLAMBなどに変更しましょう。

バッチ正規化層の直前にある畳み込み層ではバイアスを使わない

バッチ正規化層はバイアスを足し合わせるような処理が入っているため、バッチ正規化層の直前にある畳み込みそうではバイアス項は不要です3,4。

バッチ正規化層の処理イメージ

y = gamma * normalized(x) + bias  # バイアス項が考慮されている

model.zero_gradではなくparameter.grad = Noneにする

zero_grad()

  • パラメータごとにmemsetを実行する (パラメータ行列の各値に0を設定する)
  • backward実行時、parameter.grad += gradのように加算代入がおこなわれる

grad = None

  • memsetは実行されない
  • backward実行時、parameter.grad = gradのように代入がおこなわれる

backward実行時の挙動について、PyTorchの公式ドキュメントに書かれているgradientの扱いについてみてると、以下のようなことが書かれています。 parameterが現在のデータに対してbackwardを実行することで得られたgradientを受け取るとき、parameter.gradがNoneかどうかで挙動が変わります。

  • Noneのとき: parameter.grad = gradとなる
  • Noneでないとき: parameter.grad += gradになる

またparameter.grad = Noneとすることでパフォーマンスが改善する可能性があることも言及されています。

(…) is a valid alternative to model.zero_grad() or optimizer.zero_grad() that may improve performance for some networks.

デバッグ用の設定を無効にする

プログラムのデバッグ終了後は以下のデバッグ用APIの設定をFalseにしましょう

anomaly detection:

  • torch.autograd.detect_anomaly
  • torch.autograd.set_detect_anomaly(True)

autograd profiler:

  • torch.autograd.profiler.profile

automatic NVTX ranges:

  • torch.autograd.profiler.emit_nvtx

autograd gradcheck:

  • torch.autograd.gradcheck
  • torch.autograd.gradgradcheck

マルチGPUではDistributedDataParallelを使う

DataParallelはCPU1つで複数のマルチGPUを扱います。 そのためあるGPUにデータを転送している間、他のGPUは待ちの状態になってしましまいます。 さらにDataParallelはシングルノードにしか対応しないという点にも注意が必要です。

DistributedDataParallelはGPU１つにつきCPU１つが対応します。 またマルチノードにも対応しているため、容易に分散学習を適用できます。

マルチGPUでは各GPUの負荷を分散させる

マルチGPUによる学習ではbackward処理の実施後、各GPUで得られた勾配を収集するために同期処理が発生します。あるGPUがほかのGPUよりも速くforward, backwardを終えると、そのGPUはほかのGPUがforward, backwardを終えるまで待ち続けてしまいます。

たとえば自然言語処理では文長（たとえば単語数）が長いほど処理に時間がかかる傾向にあるため、各GPUが処理する文長はできるだけ均一にした方がGPUの負荷が均一になります。

apexのモジュールを使う

NVIDIAが開発しているapex (NVIDIAが開発している混合精度学習や分散学習のためのPyTorchモジュール) を使うと処理が速くなる可能性があるので検討してみましょう。

checkpointを利用して学習時のメモリ消費を減らす

通常の設定でforwardを実行すると、中間層における演算結果は、勾配の計算で再利用するため、すべて保持するようになっています5。演算結果が容易に得られる層に関しては、演算結果を保持しないようにすることで計算グラフの消費メモリを減るので、バッチサイズを大きくすることができます（forwardを再計算するよりも、バッチサイズを大きくしたほうが高速になる可能性が高い）。

通常の設定での学習

  • forward: 各層の演算結果をすべて保持するため
  • backward: forwardで計算した結果を再利用して勾配を計算する

checkpointを有効にした設定での学習

  • forward: 一部の層の演算結果だけを保持
  • backward: 演算結果が保持されていない層に関してはforwardを再計算

この機能はtorch.utils.checkpointでサポートされています。 活性化層など、再計算のコストが低い層をcheckpointを有効にして、バッチサイズを大きくすると高速化が期待できます。

torch.jit.scriptを使う

torch.jit.csriptはgeluなど要素ごとに必要な複数の演算を一つのCUDAカーネルで実行することでオーバーヘッドを減らすことができます。 たとえばgeluは以下のように、xの各要素を1.41421で割って、その結果にerf関数を適用して、…と複数の計算をします。

def gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

torch.jit.scriptを使わない場合・使う場合で以下のような違いが出てきます6。

torch.jit.scriptを使わない場合

  • カーネルを起動
  • グローバルメモリからデータxを読み込む
  • x / 1.41421を計算する
  • 計算結果をグローバルメモリに書き込む
  • グローバルメモリからデータx / 1.41421を読み込む
  • erf(x / 1.41421)を計算する
  • 計算結果をグローバルメモリに書き込む
  • …
  • 計算結果をグローバルメモリに書き込む

torch.jit.scriptを使う場合

  • カーネルを起動
  • グローバルメモリからデータxを読み込む
  • x / 1.41421を計算する
  • erf(x / 1.41421)を計算する
  • …
  • 計算結果をグローバルメモリに書き込む

上記のようにカーネルを一つにまとめることで、メモリの読み込み、書き込みのオーバーヘッドを減らすことができます。

おわり

本記事ではNVIDIAがECCV2020で発表したチュートリアル資料をもとにPyTorchプログラムを高速化するための工夫を紹介しました。 PyTorchプログラムの処理速度を改善するための機能がPyTorchに実装されており、パフォーマンスに注意が払われていることを非常に感じました。 私のような利用する側のユーザとしては非常にありがたく、また活用していきたい工夫です。

最後に、今回紹介したチュートリアルでは取り上げられなかったものの、学習時の消費メモリを削減する方法については以下で紹介しています。 もしよければこちらも御覧ください。

【PyTorch】不要になった計算グラフを削除してメモリを節約

  1. https://developer.nvidia.com/blog/how-optimize-data-transfers-cuda-cc/ ↩︎

  2. https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936/2 ↩︎

  3. https://stackoverflow.com/questions/46256747/can-not-use-both-bias-and-batch-normalization-in-convolution-layers ↩︎

  4. https://twitter.com/karpathy/status/1013245864570073090 ↩︎

  5. https://tech.preferred.jp/ja/blog/recomputation/ ↩︎

  6. https://stackoverflow.com/questions/53305830/cuda-how-does-kernel-fusion-improve-performance-on-memory-bound-applications-on ↩︎


関連記事






最近の記事