PyTorch/TensorFlowのトレーニング時にWSL2がハングする現象

結論

今回は結論のみ。

WSL2上でPyTorch/TensorFlowで画像モデルのトレーニングを行う際に、データローダーをマルチスレッドで動かすと、トレーニングがハングする現象が発生しました。

ある程度のイテレーションは学習が進むのですが、途中で突然プログレスバーが全く進まなくなるという感じの挙動です。

私のケースではWSL2がWindowsのディレクトリ(/mnt/c 以下)にアクセスする際のパフォーマンス問題によって起こった現象のようでした。なんとなくキャッシュと言うかバッファみたいなのがあって、それを超えた読み書きが行われると止まる、みたいな挙動に見えます。

類似の現象は以下のIssueにも報告されており、これを参考にLinux側のディレクトリ(/home 下など)にデータをコピーしてからトレーニングを行うことで解決しました。

github.com

github.com

おそらく最近言語モデルしかトレーニングしてなかったので問題になっていなかったのだと思うのですが、WSL2である程度大きなデータを並列読込みするようなケースでは注意が必要そうです。

環境情報

• Windows 11 Pro 22H2
• CPU: AMD Ryzen 9 7950X3D
• RAM: 64GB
• GPU: NVIDIA RTX 4090

wsl --version
WSL バージョン: 2.0.9.0
カーネル バージョン: 5.15.133.1-1
WSLg バージョン: 1.0.59
MSRDC バージョン: 1.2.4677
Direct3D バージョン: 1.611.1-81528511
DXCore バージョン: 10.0.25131.1002-220531-1700.rs-onecore-base2-hyp
Windows バージョン: 10.0.22621.2283