【PyTorch】憎き「CUDA out of memory」から学ぶGPUメモリ管理と勾配累積(Gradient Accumulation)の極意

Deep Learningエンジニアが一生のうちに最も多く目にするであろう絶望のエラーメッセージ。

「RuntimeError: CUDA out of memory. Tried to allocate xxx MiB…」

モデルをより賢くしようと意気込んで実行した数秒後、非情にもGPUのVRAM(ビデオメモリ)が溢れて強制終了するこの現象。初心者からプロのKaggle参加者まで、誰もが一度はモニタの前で頭を抱えるポイントです。

「とりあえずバッチサイズを半分にすれば動くでしょ」と思うかもしれません。もちろんそれは正しい対症療法ですが、バッチサイズを下げすぎると学習が不安定になり、本来の精度が出ないという新たな罠に陥ります。

本記事では、この憎きOOM（Out of Memory）エラーを根本から理解し、VRAMをケチりながらもバッチサイズを実質的に増やすマジック「勾配累積（Gradient Accumulation）」の仕組みや、PyTorchの高度なメモリ管理機構まで、たっぷり分かりやすく深掘りしていきます！

1. そもそもなぜVRAM（GPUメモリ）が溢れるのか？

エラーを解決する前に、まずは「何がそんなにメモリを食っているのか？」を知る必要があります。GPUのメモリをお弁当箱だと想像してみてください。モデル学習中、このお弁当箱には大きく分けて4つの巨大なおかずが詰め込まれます。

1. モデルのパラメータ（重みとバイアス）
ニューラルネットワーク自身のサイズです。しかし意外なことに、ResNetなどの標準的なモデル自体は数百MB程度と、そこまで巨大ではありません。
2. オプティマイザの保持する状態（Momentumなど）
AdamやSGDといった最適化アルゴリズムは、「過去にどの方向にどれくらい学習を進めたか」という歴史（状態）を記憶しています。実はこれがモデルパラメータの2倍〜3倍のメモリを消費します。
3. 入力データ（バッチ自体）
今回の学習で流し込む画像やテキストのデータ本体。高解像度の画像だとバッチサイズに比例して大きくなります。
4. 【真の犯人】計算グラフと中間活性化マップ（Activation Maps）
圧倒的にメモリを食いつぶしているのがこれです！PyTorchは誤差逆伝播（バックプロパゲーション）を行うため、順伝播で計算された「途中の計算結果」をすべてVRAM上に保存しています。数多くのレイヤーを通過するごとに生まれる膨大な掛け算の途中結果を記憶しているため、あっという間に数十GBのメモリが吹き飛ぶのです。

つまり、「画像が大きく（3）、ネットワークが深い（4）」ほど、保存しておくべき途中結果が爆発的に増え、すぐにお弁当箱が溢れてしまう（OOMエラー）というわけです。

2. 「勾配累積 (Gradient Accumulation)」という究極のズル

バッチサイズを大きくしたい（例えば64にしたい）けれど、VRAMの限界でバッチサイズ16しか学習機に乗らない。そんな絶望的な状況を救ってくれるのが「勾配累積（Gradient Accumulation）」です。

仕組みはとてもシンプルで賢いものです。
1. バッチサイズ16で順伝播と逆伝播を行う。
2. ここで重みを「更新しない」。計算された勾配（エラーの方向）だけを一時的にメモ（加算）しておく。
3. この作業を4回繰り返す。（16 × 4回 = 合計64個のデータを見たことになる）
4. 4回分の勾配が溜まったところで、初めて1回だけ重みを更新する！

このようにすることで、VRAMの消費量はバッチサイズ16のままなのに、数学的には「バッチサイズ64で学習した」のとほぼ同等の安定した学習効果を得ることができるのです。

勾配累積の具体的な実装コード

PyTorchでの実装は非常に簡単です。以下のようにループの中に数行足すだけで完成します。

# 勾配を累積するステップ数を定義
# 実質バッチサイズ = 実際のバッチサイズ(16) × accumulation_steps(4) = 64
accumulation_steps = 4
# 【重要】ゼロクリアはミニバッチごとではなく、ループの外に出す
optimizer.zero_grad() 
for i, (inputs, labels) in enumerate(dataloader):
    inputs, labels = inputs.to(device), labels.to(device)
    
    # 通常の順伝播とロス計算
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # 【重要】ロスを accumulation_steps で割る
    # 勾配を足し合わせるので、スケールを揃えるために平均化します
    loss = loss / accumulation_steps
    
    # 勾配の計算 (ここではまだoptimizer.step()を呼ばないため、勾配が「加算」され続ける)
    loss.backward()
    
    # 指定ステップ数（今回は4回ごと）に達したときだけ、実際に重みを更新する！
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad() # ここでようやく次の累積に向けてゼロクリア

3. さらなるメモリ節約の切り札：自動混合精度（AMP）

勾配累積でも厳しい場合、もう一つの強力なメモリ削減手法が PyTorch AMP (Automatic Mixed Precision) です。

AIの計算の大部分は、実はそこまで細かい「小数点以下の精度」を必要としません。
通常32ビット（Float32）で計算しているテンソルを、影響の少ない部分だけ16ビット（Float16）に自動で落として計算させる機能がAMPです。これにより、なんとVRAM消費量がほぼ半減し、さらに最近のGPU（Tensorコア搭載）であれば計算速度が最大で2〜3倍に跳ね上がるという「使わない手はない」最強の機能です。

from torch.cuda.amp import autocast, GradScaler
# 勾配ロススケールを管理するオブジェクト
scaler = GradScaler()
for inputs, labels in dataloader:
    optimizer.zero_grad()
    
    # この with 文の中の計算だけが、自動的にFP16になりメモリが半減！
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
    # ロスをスケーリングしてバックプロパゲーション
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()