小ネタ：Pytorch で Automatic Mixed Precision (AMP) の ON/OFF をするときの話 - 俵言

最近気付いたのでメモ。長くなってしまったので、結論だけ見たい場合はまとめまで読み飛ばしてください。 まえおき NN を学習する際の GPU メモリの使用量軽減や学習速度の向上手段として混合精度学習(Mixed Precision Training) は欠かせません。pytorch では torch.cuda.amp モジュールを用いることでとてもお手軽に使うことが可能です。 以下は official docs に Typical Mixed Precision Training と題して載っている例ですが 、 model の forward と loss の計算を amp.autocast の with 文中で行い、loss の backward と optimizer の step に amp.GradScaler を介在させています*1。 # Creates model and opt

[sh19910711]

"混合精度学習: ちょっとの変更でGPU消費量が半分くらいになり計算速度も2倍ぐらいになる / 何も考えずおまじない的に autocast と GradScaler を書いていたのですがもっと早く docs を読むべきでした" 2021