pytorchのnon_blockingという機能が使えると思った
この投稿はrioyokotalab Advent Calendar 2020 19日目の投稿です。
デバイス転送部分がボトルネックになる可能性
PyTorchのCUDAで動作する部分は非同期で動いてくれるらしいです。[要出典]なので、特に何も意識しなくても、データの読み出しとかにかかるCPUの計算時間をモデルのforward, backwardの計算時間に隠してくれるようにできています。つまり、CPU側の作業とGPU側の作業を並行して行ってくれるということです。
この時、データをCPUメモリからGPUメモリに移動する作業が時間がかかる場合があります。そこを、非同期にするオプションがあるらしいです。
CPUメモリ上のtensor
をGPUメモリに移すためには、to
メソッドを呼びます。このto
メソッドの引数にnon_blocking
という引数があり、これにTrue
を渡すことで、非同期に転送ができるらしいです。
non_blockingについての議論されているのはこちら discuss.pytorch.org
non_blockingで速度が上がったという報告はこちら qiita.com
実験
ひなどりクラスタ上で試してみました。確保したノードはtitanv
です
とりあえず、
- 画像サイズ512x512
- バッチサイズ 64
- モデル ResNet18
- DataLoader num_worker=1, pin_memory=True
で試してみました。
non_blocking=True
non_blocking=False
うーん... 別段早くはならない...
まとめ
いかがでしたか?non_blockingについて試してみましたが別段早くはなりませんでした。
載せた実験以外にもいくつか試してみましたが、明確に早くなったものはありませんでした。
必要な条件がわかったら追記しようと思います。