【解決】 CUDA out of memory の解決方法と原因 | PyTorch/TensorFlow トラブルシューティング

PyTorchやTensorFlowで深層学習モデルを開発・実行している際に「CUDA out of memory」というエラーに遭遇すると、非常に困惑し、作業が中断されてしまいますよね。しかし、ご安心ください。このエラーはGPUメモリが不足していることを示すもので、多くの場合、比較的簡単な手順で解決できます。

この記事では、Windowsユーザー向けに、このエラーの概要から、今すぐ試せる具体的な解決策、さらには恒久的な再発防止策までを、ロジカルかつ分かりやすく解説します。

1. CUDA out of memory とは?(概要と緊急度)

CUDA out of memory エラーは、一言で言えば「GPU(グラフィックス処理ユニット)のメモリが足りません!」という悲鳴です。

深層学習モデルの学習や推論では、大量の数値計算がGPU上で行われます。このとき、モデルのパラメータ、入力データ、中間計算結果などがGPUメモリ上に展開されます。もし、これらのデータがGPUの搭載メモリ容量を超えてしまうと、このエラーが発生し、プログラムの実行が強制的に停止してしまいます。

このエラーが発生すると、それ以上学習や推論を進めることができないため、即座の対処が必要です。</

2. 【最速】今すぐ試すべき解決策

まずは、最も効果的で簡単に試せる解決策からご紹介します。これらの手順を試すことで、多くの「CUDA out of memory」エラーは解決します。

解決策1:バッチサイズを小さくする

最も簡単で、かつ非常に効果的な方法です。バッチサイズとは、一度にGPUに読み込ませて処理するデータの数を指します。これを小さくすることで、一度にGPUが使用するメモリ量を直接的に削減できます。

通常、モデルの学習コード内でbatch_sizeとして定義されている箇所を探し、値を小さく変更してください。例えば、32から168、あるいは4に減らしてみましょう。

PyTorchでの例:

# 変更前:
# dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)

# 変更後:
dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True) # バッチサイズを半分に
    

TensorFlowでの例:

# 変更前:
# model.fit(train_dataset, batch_size=32, epochs=10)

# 変更後:
model.fit(train_dataset, batch_size=16, epochs=10) # バッチサイズを半分に
    

解決策2:不要なGPUメモリを解放する

学習中に過去のイテレーションで生成された中間結果や、デバッグ用に作成されたテンソルなどがGPUメモリ上に残り続けている場合があります。これらを明示的に解放することで、メモリを確保できます。

特にPyTorchでは、明示的なキャッシュクリアが有効な場合があります。

PyTorchでの例:

import torch

# 不要になった変数(テンソル)を削除
del some_large_tensor
del another_temp_variable

# GPUキャッシュを明示的にクリア(非常に重要!)
torch.cuda.empty_cache()
    

TensorFlowでは、通常はガベージコレクションによって自動的に管理されますが、tf.config.experimental.set_memory_growth(gpu, True)を使用することで、必要な分だけメモリを割り当てるように設定できます。これは起動時に一度だけ実行します。

TensorFlowでの例 (起動時の設定):

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # GPUメモリの使用量を必要に応じて増やす設定
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # メモリ成長を設定する際は、プログラムの開始時に行う必要があります
        print(e)
    

解決策3:他のGPUプロセスを特定・終了する (Windowsユーザー向け)

あなたのPythonスクリプト以外にも、他のアプリケーションやプロセスがGPUメモリを使用している可能性があります。例えば、別の学習スクリプト、ゲーム、動画編集ソフトウェア、あるいはOSのGUIなどが原因かもしれません。

以下のコマンドを使って、現在GPUを使用しているプロセスを確認し、不要なものを終了しましょう。

Step 1: GPUの使用状況を確認する

PowerShellまたはコマンドプロンプトを開き、nvidia-smiコマンドを実行します。

nvidia-smi
    

このコマンドを実行すると、現在GPU上で動作しているプロセスの一覧と、それぞれのプロセスが使用しているGPUメモリの量が表示されます。「Processes」セクションを確認し、見覚えのないプロセスや、不要なプロセスが大量のメモリを消費していないか確認してください。

例えば、以下のような出力が得られます。(実際の出力はもっと詳細です)

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   PID   Type   Process name                             GPU Memory   |
|        ID                                                      Usage      |
|=============================================================================|
|    0   1234  C      python.exe                                 2000MiB    |
|    0   5678  C      chrome.exe                                  500MiB    |
|    0   9101  C      another_script.exe                         4000MiB    |
+-----------------------------------------------------------------------------+
    

Step 2: 不要なプロセスを終了する

nvidia-smiで特定した不要なプロセスのPID(Process ID)をメモします。

タスクマネージャーを使用する場合:

  1. Ctrl + Shift + Escキーを押してタスクマネージャーを開きます。
  2. 「詳細」タブに移動します。
  3. 先ほどメモしたPIDと一致するプロセスを探し、右クリックして「タスクの終了」を選択します。

コマンドプロンプト(またはPowerShell)を使用する場合:

特定したPIDを持つプロセスを強制終了するコマンドです。誤って必要なプロセスを終了しないよう、PIDをよく確認してください。

taskkill /PID [プロセスのID] /F
    

例: PIDが9101のプロセスを終了する場合

taskkill /PID 9101 /F
    

これにより、そのプロセスが使用していたGPUメモリが解放され、あなたのスクリプトが利用できるようになります。

3. CUDA out of memory が発生する主要な原因(複数)

上記の即効性のある解決策で問題が解決しない、あるいは再発する場合は、根本的な原因を理解し、より恒久的な対策を検討する必要があります。主な原因は以下の通りです。

原因1:バッチサイズが大きすぎる

前述の通り、一度に処理するデータ量が多すぎると、GPUメモリが不足します。特に、高解像度の画像や長いシーケンスデータを使用する場合、バッチサイズを小さくする必要があるかもしれません。

原因2:モデルの複雑さや入力データのサイズ

  • モデルのパラメータ数が多い: 非常に深いネットワークや、幅広い層を持つモデルは、多くのパラメータを持ち、それらがGPUメモリを消費します。
  • 入力データの次元が高い: 高解像度の画像、長いテキストシーケンス、3Dデータなどは、それ自体が大量のメモリを占有します。
  • 中間特徴マップが大きい: CNNの層を深くしていくと、中間特徴マップのサイズが大きくなり、メモリ消費が増えることがあります。

原因3:GPUメモリリークまたは不要なテンソル

PythonのガベージコレクタはCPUメモリを管理しますが、GPUメモリの管理はフレームワーク(PyTorch/TensorFlow)に依存します。明示的に変数を削除しないと、GPUメモリが解放されずに残り続ける「メモリリーク」のような状態になることがあります。

  • 学習ループ内でテンソルを繰り返し作成し、deltorch.cuda.empty_cache()を呼び忘れている。
  • デバッグ目的で作成した大きなテンソルを削除し忘れている。
  • 勾配計算が不要な部分でtorch.no_grad()(PyTorch)やtf.GradientTape(persistent=False)(TensorFlow)を使用し忘れている(勾配計算のためのグラフ構築もメモリを消費します)。

原因4:他のアプリケーションやOSプロセスによるGPUメモリ消費

前述の通り、GPUは深層学習のためだけに存在するわけではありません。Windows OSのグラフィカルインターフェース、Webブラウザ、ゲーム、他のGPUを利用するアプリケーションなどが、知らず知らずのうちにGPUメモリを消費していることがあります。

4. PyTorch/TensorFlowで恒久的に再発を防ぐには

CUDA out of memory」エラーの再発を防ぐためには、より戦略的なアプローチが必要です。以下の方法を検討してみましょう。

対策1:メモリ使用量の継続的な監視

開発中に定期的にGPUメモリの使用状況をチェックする習慣をつけましょう。

  • nvidia-smiコマンド: PowerShellやコマンドプロンプトで実行し、GPUの現在の使用状況を監視します。
    nvidia-smi -l 1 # 1秒ごとに更新して表示
  • PyTorchのメモリプロファイラ:
    • torch.cuda.memory_allocated(): 現在割り当てられているGPUメモリ量
    • torch.cuda.max_memory_allocated(): ピーク時のGPUメモリ割り当て量
    • torch.cuda.memory_summary(): 詳細なメモリレポート
  • TensorFlowのメモリ設定: tf.config.experimental.set_memory_growth(gpu, True)をプログラムの開始時に設定し、必要な分だけGPUメモリを確保するようにします。

対策2:混合精度トレーニング (Mixed Precision Training) の導入

多くの深層学習モデルは、計算に32ビット浮動小数点数(FP32)を使用します。しかし、ほとんどの最新GPUは16ビット浮動小数点数(FP16)での計算もサポートしており、これを利用することでメモリ使用量を約半分に削減し、計算速度も向上させることができます。

PyTorchでの例 (torch.cuda.ampを使用):

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for input, target in data:
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    

TensorFlowでの例 (tf.keras.mixed_precisionを使用):

from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy('mixed_float16')

# モデル構築や学習は通常通り行います。
# Kerasが自動的に混合精度を活用します。
model = tf.keras.models.Sequential([ ... ])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(train_dataset, epochs=10)
    

対策3:勾配チェックポインティング (Gradient Checkpointing) の活用

非常に大きなモデル(特にTransformerのようなモデル)を扱う場合、すべての層の中間出力をメモリに保持すると、すぐにメモリが枯渇してしまいます。

勾配チェックポインティングは、フォワードパスの一部の中間結果のみを保存し、バックワードパスで必要になったときに再計算することで、メモリ消費を大幅に削減する手法です。計算時間は少し増えますが、メモリ効率が向上します。

PyTorchでの例 (torch.utils.checkpointを使用):

import torch.utils.checkpoint as checkpoint

# モデルの一部 (例: 一つのレイヤーグループ) をcheckpoint化
def model_segment(x):
    # この中で複数のレイヤーを定義
    return some_layer(x)

# 学習ループ内で
output = checkpoint.checkpoint(model_segment, input_tensor)
    

対策4:モデル構造やデータパイプラインの最適化

  • モデルの枝刈り(Pruning)や量子化(Quantization): モデルのサイズを小さくしたり、より低精度な数値表現に変換したりすることで、メモリと計算リソースを節約します。
  • データローダーの最適化: PyTorchのnum_workerspin_memory=Trueオプションを適切に設定することで、データ転送の効率を上げ、CPUとGPU間のボトルネックを解消できる場合があります。ただし、num_workersを増やしすぎるとCPUメモリを消費することもあるため、適切なバランスを見つけることが重要です。

対策5:より大容量のGPUへのアップグレード (最終手段)

上記のソフトウェア的な解決策を試してもなおメモリ不足が解消しない場合、あるいは大規模なモデルやデータを扱う必要がある場合は、より多くのGPUメモリを搭載したグラフィックカードへのアップグレードを検討する時期かもしれません。これは最もコストがかかる解決策ですが、最も根本的な解決策でもあります。

CUDA out of memory」エラーは、深層学習における一般的な課題の一つです。この記事で紹介した解決策を一つずつ試すことで、きっとあなたの問題も解決し、スムーズな開発・研究活動に戻ることができるでしょう。