pytorchのdata load時に処理が止まる

dataloader経由でデータを取得する際に処理待ちのような状態になって動かなくなることがあったのでメモ。

具体的にやろうとしてたのはDETRを動かしてみることで、画像2枚だけの小さなデータセットを用意して実行したところ以下の部分以降に処理が進まなかった
https://github.com/facebookresearch/detr/blob/eb9f7e03ed8e2ed2cd55528989fe7df890bc3fc0/datasets/coco.py#L29

• 前後にprintを置いて後ろの出力だけが出ない
• エラーが出るわけではなく学習処理は実行中のまま進まない
• 上記のtransformsを単純に ToTensor() だけに置き換えてみても同様

学習開始の際に以下のようなwarningが出てたのでそれも調べてみる

[W ParallelNative.cpp:212] Warning: Cannot set number of intraop threads after parallel work has started or after set_num_threads call when using native parallel backend (functio
n set_num_threads)

メッセージでググるとこういうのがヒットする https://stackoverflow.com/questions/64772335/pytorch-w-parallelnative-cpp206 https://github.com/pytorch/pytorch/issues/46409

私の環境は

• macOS
• python3.7.7
• pytorch1.9.0

なのでリンクにある状態とは異なるものの、ひとまず num_worker=0 として試したらwarningは出なくなり、かつ処理も進むようになった。というわけでdataloaderが複数workerで動いていることによる問題だった模様。

OMP_NUM_THREADS=1, seems to resolve the issue. The num_workers parameters on DataLoader seems to have no effect on the issue.

githubのissueだとこのように書いてあるが、pytorchのversionが異なるためかnum_workerの設定だけで問題なかった。

ちなみにせっかくなので調べてみたところ OMP_NUM_THREADS はCPU上での並列処理のスレッド数を設定するものだった
https://pytorch.org/docs/stable/notes/cpu_threading_torchscript_inference.html https://qiita.com/sakaia/items/6d31dedb396463023cfe

この設定にすると当然dataloaderのworkerの並列動作がなくなるのでパフォーマンスは低下するが、今回はまず動かしてみるという状態なので今の所はこれで解決。