multi-processのDataLoaderにおけるメモリ消費の問題について

問題

multi-processでPytorchのデータローダを実行する場合、single-processの場合と比べてメモリ消費が増大することがある。[1]ではその原因と対処方法について説明されている。

原因

pytorchのデータローダで大量のpythonオブジェクト(list, dict etc.)を参照する場合に、オブジェクトの参照カウントがプロセス固有のメモリ空間に複製され、同じ内容のオブジェクトをメモリを重複して保持する。したがって、dataset内に大量のオブジェクトを保持するとメモリ消費を圧迫する原因となりうる。

対策

pythonオブジェクトの数を減らす。具体的にはdictやlistの代わりにnumpy配列やPytorchのテンソル、pandasのdataframeの形式で保持する。

Reference

[1] Demystify RAM Usage in Multi-Process Data Loaders