【PyTorch】Gradientのnan/inf発生時のデバッグ法

PyTorchのGradient nan/inf問題

PyTorchの使用中にGradient nan/inf問題に遭遇する場合があります。本記事ではこの問題の原因と解決方法を詳しく解説します。

エラーの症状

学習中にlossがNaNやInfになる、またはgradientにNaN値が含まれてモデルが収束しない問題です。

原因

1. Learning Rateが大きすぎる
勾配爆発の最も一般的な原因です。

2. 数値的に不安定な演算
log(0)、div(0)、exp(大きな値)などが原因です。

3. データの異常値
入力データにNaNやInfが含まれている場合です。

解決方法

手順1: anomaly detectionの有効化

# NaN発生箇所を特定
torch.autograd.set_detect_anomaly(True)

# 学習ループ内で自動検出
with torch.autograd.detect_anomaly():
    output = model(input)
    loss = criterion(output, target)
    loss.backward()  # NaN発生時にスタックトレースが表示される

手順2: Gradient Clipping

# 勾配の最大ノルムを制限
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# または値ベースのクリッピング
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

手順3: 入力データの検証

# データにNaN/Infがないか確認
assert not torch.isnan(input_data).any(), "入力にNaNが含まれています"
assert not torch.isinf(input_data).any(), "入力にInfが含まれています"

まとめ

本記事で紹介した手順を試すことで、多くの場合問題を解決できます。環境固有の問題が残る場合は、公式ドキュメントやコミュニティフォーラムを参照してください。