pytorchのnon_blockingという機能が使えると思った

この投稿はrioyokotalab Advent Calendar 2020 19日目の投稿です。

adventar.org

バイス転送部分がボトルネックになる可能性

PyTorchのCUDAで動作する部分は非同期で動いてくれるらしいです。[要出典]なので、特に何も意識しなくても、データの読み出しとかにかかるCPUの計算時間をモデルのforward, backwardの計算時間に隠してくれるようにできています。つまり、CPU側の作業とGPU側の作業を並行して行ってくれるということです。

この時、データをCPUメモリからGPUメモリに移動する作業が時間がかかる場合があります。そこを、非同期にするオプションがあるらしいです。

CPUメモリ上のtensorGPUメモリに移すためには、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

f:id:deoxy:20201219234107p:plain

non_blocking=False

f:id:deoxy:20201219234144p:plain

うーん... 別段早くはならない...

まとめ

いかがでしたか?non_blockingについて試してみましたが別段早くはなりませんでした。

載せた実験以外にもいくつか試してみましたが、明確に早くなったものはありませんでした。

必要な条件がわかったら追記しようと思います。