Tech

PyTorchの「CUDA out of memory」の原因と対処

学習や推論を回すと、PyTorchがこう言って止まる。

RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB ...

GPUのメモリが足りていない。原因はバッチサイズが大きい、モデルが大きい、推論なのに勾配を保持している、メモリの断片化、などだ。効く順に対処する。

まずバッチサイズを下げる

いちばん確実で効く。半分にして通るか試す。

batch_size = 16  # 32 -> 16 のように下げる

推論では勾配を持たない

評価や推論で勾配を保持すると、無駄にメモリを食う。

model.eval()
with torch.no_grad():
    out = model(x)

学習はAMP（混合精度）と勾配累積

混合精度で使用メモリが減る。さらに、小さいバッチを複数回ためてから更新すれば、実効バッチサイズを保ちつつメモリを抑えられる。

scaler = torch.cuda.amp.GradScaler()
for x, y in loader:
    with torch.cuda.amp.autocast():
        loss = criterion(model(x), y)
    scaler.scale(loss).backward()
    scaler.step(optimizer); scaler.update(); optimizer.zero_grad()

断片化が原因のとき

空きはあるのに確保できない場合は断片化が疑わしい。次の環境変数が効くことがある。

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

まとめ

  • まずバッチサイズを下げる
  • 推論は model.eval() と torch.no_grad()
  • 学習はAMPと勾配累積でメモリを節約
  • 「空きはあるのに落ちる」なら断片化対策の環境変数を試す