【PyTorch】Mixed Precision (AMP) 学習のエラー解決法

PyTorchで発生するMixed Precision学習エラー

PyTorchの使用中にMixed Precision学習エラーが発生することがあります。本記事では原因と具体的な解決方法を解説します。

エラーの症状

torch.cuda.amp（Automatic Mixed Precision）を使用した学習で、lossがNaNになる、GradScalerのオーバーフロー、またはモデルの収束が悪化する問題です。

原因

1. Loss Scaleのオーバーフロー
FP16の表現範囲を超えるgradient値が発生するとNaNになります。

2. amp非対応の演算
一部のカスタム演算やloss関数がFP16に対応していない場合があります。

3. Learning Rateの不適合
FP32で最適だった学習率がAMP環境では大きすぎる場合があります。

解決方法

手順1: 正しいAMPの実装パターン

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(batch)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

手順2: NaN対策

# Gradient Clippingを追加
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)

手順3: 特定の演算をFP32に固定

# Loss計算をFP32で実行
with autocast():
    output = model(input)

# autocastの外でloss計算
loss = criterion(output.float(), target)

まとめ

PyTorchのMixed Precision学習エラーについて解説しました。環境設定の見直しとバージョン確認を最初に行うことが、最も効率的なトラブルシューティングの第一歩です。