[GPUを簡単に高速化・省メモリ化] NVIDIAのapex.ampがPyTorchに統合されたようです

#対象とする読者
PyTorch 1.6を利用できる環境にあってVolta以降のGPUの計算を簡単に高速化・省メモリ化したい方

#はじめに

Volta以降のGPUではfloat16/float32混合精度演算により計算の高速化・メモリの節約を行うTensorCoreが搭載されています。

しかし、特に明示的にプログラムを書かない限り?はfloat32による計算が行われ、TensorCoreを有効活用することはできません。かといって単に全てfloat16化して計算を行うだけだと、表現力が少なかったり、勾配でアンダーフローしてしまったりとモデルの性能が維持できません。

そこで、Automatic Mixed Precision (amp)機能によってfloat32の保持や勾配の値のスケーリング等を自動的に行うことでTensorCoreを有効活用し、GPU計算の高速化・省メモリ化が実現されます。

もう少ししっかり知りたい方は以下の記事・論文を参照してみてください。

• MIXED PRECISION TRAINING (ICLR 2018)

• Pytorchなら同じGPUで2倍早い学習ができるっていったんだよ！

• Automatic Mixed Precision (amp) でニューラルネットワークのトレーニングを高速化

PyTorchにおけるampの利用

PyTorchにおいてはamp機能としてNVIDIA社のapexがよく利用されていると思います。しかし、環境によってはいくつかのライブラリのversionを上手く対応させたり、gitのcommitを少し戻したりしないと上手く動かないことがあります。（僕はインストール時に沼にハマりました。）

PyTorch 1.6からはこのapex.ampがPyTorchに統合されています。apex.ampにあった問題が何点か解消されてより簡単に導入できるようになりました。公式でもPyTorchに統合した方への移行が推奨されています。

#使い方

import torch 
# スケーラーを定義
scaler = torch.cuda.amp.GradScaler() 

for data, label in data_iter: 
   optimizer.zero_grad() 

   # 演算処理をampの対象にキャスト
   with torch.cuda.amp.autocast(): 
      loss = model(data) 

   # スケールした勾配を作る
   scaler.scale(loss).backward() 

   # 勾配をアンスケールしてパラメータの更新
   scaler.step(optimizer) 

   # スケーラーの更新
   scaler.update() 

おわりに

見ていただきありがとうございます。
怪しい箇所があれば、お気軽にコメントしていただけると幸いです。

参考

Introducing native PyTorch automatic mixed precision for faster training on NVIDIA GPUs