PyTorchで発生するCUDA out of memoryエラーの完全な解決ガイド

冒頭：どんな問題が発生するか

PyTorchでGPUを使用して機械学習モデルを実行していると、唐次に以下のようなエラーに遭遇することがあります。

RuntimeError: CUDA out of memory. Tried to allocate 70.00 MiB (GPU 0; 4.00 GiB total capacity; 3.50 GiB already allocated; 363.19 MiB free; 6.46 GiB reserved in total by PyTorch)

または

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.60 GiB (GPU 0; 23.65 GiB total capacity; 14.12 GiB already allocated; 3.97 GiB free; 18.15 GiB reserved in total by PyTorch)

このエラーは、GPUのVRAM（ビデオメモリ）が不足していることを示しています。エラーログには「total capacity」（総容量）、「already allocated」（既に使用中）、「free」（空き容量）、「reserved」（PyTorchが予約済みのメモリ）が表示されるため、これらの数値を確認することが問題解決の第一步となります。

結論：解決策の要約

CUDA out of memoryエラーは、①メモリの適切な開放、②バッチサイズの縮小、③メモリ最適化技術の適用、④ 不要なプロセスの終了の4つの観点から対処することで大半の場合解決可能です。以下では、具体的な対処法をステップバイステップで解説します。

具体的な手順：ステップバイステップでの解決法

ステップ1：現在のGPUメモリ使用状況を確認する

まず、現在どのような状態でGPUメモリが使用されているかを把握しましょう。

import torch

# GPU情報の確認
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA version: {torch.version.cuda}")
print(f"GPU count: {torch.cuda.device_count()}")

# 現在使用中のGPUメモリを確認
if torch.cuda.is_available():
    print(f"nCurrent GPU: {torch.cuda.current_device()}")
    print(f"GPU name: {torch.cuda.get_device_name(0)}")
    print(f"Allocated memory: {torch.cuda.memory_allocated(0) / 1024**3:.2f} GB")
    print(f"Reserved memory: {torch.cuda.memory_reserved(0) / 1024**3:.2f} GB")
    print(f"Max allocated memory: {torch.cuda.max_memory_allocated(0) / 1024**3:.2f} GB")

ステップ2：メモリをクリアする（最も簡単な対策）

学習済みモデルやテンソルがキャッシュとしてメモリに残ったままになっている場合があります。

# 方法1：キャッシュをクリア
torch.cuda.empty_cache()

# 方法2：すべてのGPUメモリを解放
torch.cuda.synchronize()
del 変数名  # 不要な変数を削除
torch.cuda.empty_cache()

# 方法3：Pythonのガベージコレクションを強制的に実行
import gc
gc.collect()
torch.cuda.empty_cache()

ステップ3：バッチサイズの縮小

最も効果的な対策の一つは、バッチサイズ（batch_size）を小さくすることです。

# 元のコード
# batch_size = 32

# 修正後：バッチサイズを半分に削減
batch_size = 16  # または8, 4など，逐步的に小さくして試す

# データローダーの更新
train_loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=4
)

ステップ4：データ型の最適化（FP32からFP16へ）

メモリ使用量を半分に抑えるために、FP16（半精度浮動小数点）を使用します。

# FP16（半精度）への変換
model = model.half()  # モデル全体を半精度に変換
input_tensor = input_tensor.half()  # 入力テンソルも半精度に

# PyTorch 1.6以降ではautocastを使用可能
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in train_loader:
    data, target = data.cuda(), target.cuda()
    
    optimizer.zero_grad()
    
    # autocastで自動的にFP16に変換
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    # 勾配計算をFP32で行う（安定性のため）
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

ステップ5：不要になった変数の明示的な削除

ループ内で新しい変数を作成し続けた場合、メモリリークが発生しやすくなります。

# 悪い例：ループ内で変数が蓄積
for data in dataloader:
    output = model(data)
    loss = criterion(output, target)
    # output, lossが蓄積されていく

# 良い例：明示的に変数を削除
for data in dataloader:
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    
    # 明示的なメモリ開放
    del output, loss
    torch.cuda.empty_cache()

ステップ6：Gradient Checkpointing（勾配チェックポインティング）の適用

メモリオーバーヘッドを削減するために、勾配の計算方法を変更します。

# モデル定義時にgradient checkpointingを有効化
from torch.utils.checkpoint import checkpoint_sequential

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )
    
    def forward(self, x):
        # 層をグループ化することでメモリを節約
        modules = [self.layers[i] for i in range(len(self.layers))]
        return checkpoint_sequential(modules, 2, x)

ステップ7：デッドプロセスの確認と終了

ゾンビ化したPythonプロセスがGPUメモリを使用している場合があります。

# Linux/Macの場合：GPUを使用中のプロセスを確認
!nvidia-smi

# Pythonプロセスを確認
!ps aux | grep python

# 特定のプロセスを強制終了
# !kill -9 [PID]

Windowsの場合は、タスクマネージャーからpython.exeプロセスを探して終了させてください。

ステップ8：新しいPythonセッションで再実行

CUDAコンテキストが破損している場合があります。その場合は現在のPythonプロセスを終了し、新しいセッションで実行してください。

補足・注意点

バージョン依存の注意

PyTorch 1.6以降：自動混合精度（AMP）が正式サポートされています
PyTorch 1.9以降：memory_stats()メソッドで詳細なメモリ使用統計を取得可能
CUDA 11.0以上：より効率的なメモリ管理が実装されています

環境別の注意点

Windows環境：GPU驱动 программаが正しくインストールされているか確認が必要です
Linux環境：nvidia-smiコマンドで詳細なGPU情報が取得可能
クラウド環境（AWS/GCP）：インスタンスタイプによって使用可能なVRAM量が異なるため、適切なインスタンスを選択してください

よくある落とし穴

モデルの可視化でメモリを消費：tensorboardやwandb使用時に追加のメモリが必要
DataLoaderのnum_workers：多く設定しすぎるとメモリ使用量が増加
validationフェーズでもdropoutが有効：model.eval()忘れで余分なメモリを使用
マルチGPU使用時のメモリ分散：DataParallel使用時は各GPUにモデルがコピーされる

【PyTorch】CUDA out of memoryエラーの解決法と防止策