Pytorch Torch OUT OF MEMORY - GPUメモリ不足の対処法

以下はディープラーニングの作業中によく遭遇するGPUメモリ不足への対処法です。

1. 不要な変数のクリア

PyTorchのメモリキャッシュのクリア

import torch
torch.cuda.empty_cache()

使用していない変数によって占有されているメモリへの参照をクリアします。

2. Pythonのガベージコレクション

Pythonのガベージコレクタを使用して変数を削除し、メモリを解放

import gc
gc.collect()

3. バッチサイズの調整

カーネルを再起動した後、バッチサイズを小さくして最適なサイズを見つける

4. ネットワーク構造の調整

• ネットワークのレイヤーのサイズを変更
• レイヤーを分割

5. GPUメモリ割り当ての理解

torch.cuda.memory_summary(device=None, abbreviated=False)

これでメモリ割り当ての読みやすい要約が得られます。