PyTorch 1.6の新機能の話

先日より色々と深層学習で試行錯誤をしているとミニバッチサイズをパラメータとして振りたくなったが、メモリ不足で上限があまり上げられなかった。

たまたま、相談したところV100搭載機を借りることができたのでGPUメモリが32GBの環境が使えることになった。

 

環境を整えるところでPyTorchの最新版が1.6になっていることに気づく。

この手のは今まで使っていた旧バージョンとの差異で躓くリスクと新機能の恩恵をなんとなくイメージして迷うところだけど、深層学習分野は劇的に進化中なので即最新の1.6をインストールする。

 

普段はWindows環境なのだがお借りしたのがUbuntuだったので環境の差異に戸惑っていたが2,3日で何とか同じプログラムを動かせるようになった。予定していたミニバッチサイズも簡単に今までの二倍以上はいけそうだ。

 

ところで、以下のような話題も気になっておりそのうちテストしたいなぁと思っていた。

qiita.com

 

具体的には計算精度を単精度から半精度などに落とすことにより演算速度を二倍、使用メモリを半分にするような機能である。NVIDIAGPUだと最近のやつは使えるそうだ。

ただ、ライブラリとして出たばかりで完成度に不安があった。

 

ところが、本件の仕組みをPyTorch 1.6では公式に取り込まれていた上に、Windows環境も動作対象に入っていた。

ってことで公式ドキュメントを元にやってみたところ、わずか10分足らずのコーディングで二倍速を得ることができた。使用メモリは半分までにはならなかった。

 

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

pytorch.org

 

とても簡単で性能が上がるので是非試していただきたい。

確認事項はGPUとしてTensorCore搭載、PyTorchは1.6です。OSはWindowsでも動きます。

---

追記:

元ネタはこの辺だからとても速い対応で公式に取り込まれたんだとおもう。

github.com

---

追記:

本機能のテストとしてdlshogiの学習を高速化しておいたものを本家に送っておいた。無事より良い形で取り入れられたようだ。

---

追記2:

うちでも測定していたが、やはり二倍速くらいだそうな。

 

tadaoyamaoka.hatenablog.com