잘못 행동하는 기울기 디버깅

역전파가 잘못될 때 쓰는 도구

학습이 발산하거나 손실이 NaN이 되거나 기울기가 까닭 없이 0 또는 무한대가 되는 건 전형적인 이상 징후야. 다음 네 가지 PyTorch 도구만 알아도 대부분의 문제를 진단할 수 있어:

기울기 노름을 기록해. 역전파 뒤에 clip_grad_norm_(model.parameters(), float('inf'))를 추가하고 반환값을 기록해. 값이 갑자기 1e6까지 튄다면 바로 앞 단계에서 무언가 잘못됐다는 신호야.
NaN과 Inf를 감지해. 매 단계에서 torch.isfinite(loss)를 확인하면 적은 비용으로 Inf와 NaN을 모두 잡을 수 있어. torch.autograd.set_detect_anomaly(True)는 NaN 역전파를 만든 정확한 연산을 알려 줘. 느리지만 원인을 찾기 어려울 때 아주 유용해. 다만 범위를 주의해: 이상 탐지는 NaN에만 반응하고 Inf에는 반응하지 않아. 아래 'Inf와 NaN' 절에서 차이를 살펴볼게.
매개변수별 기울기 통계를 봐. 역전파 뒤에 model.named_parameters()를 순회하며 평균, 표준편차, 최댓값을 출력해. 기울기 소실은 이 통계에서 금방 드러나.
모양 불일치는 기울기 훅으로 확인해. 계층 출력에 훅을 등록하고 기울기 모양이 예상과 같은지 검증해. 건너뛰기 연결이나 사용자 정의 계층을 추가할 때 특히 유용해.

Inf와 NaN: 이상 탐지의 범위는 좁아

많은 튜토리얼이 놓치는 부분이 있어. set_detect_anomaly(True)는 역전파 출력이 NaN일 때만 오류를 내고 Inf에는 반응하지 않아. 같은 함수도 입력에 따라 둘 중 하나를 만들 수 있어.

Inf가 생기는 경우: 이상 탐지는 반응하지 않으므로 torch.isfinite()로 잡아:

sqrt(0).backward(): x=0에서 도함수 1/(2·sqrt(x))가 ∞가 돼. (x + 1e-9).sqrt()처럼 작은 값을 더해 해결해.
log(0).backward(): x=0에서 도함수 1/x가 ∞가 돼. (p + 1e-9).log()처럼 작은 값을 더해 해결해.

NaN이 생기는 경우: 이상 탐지가 작동해 정확한 연산을 지목해 줘:

sqrt(-1), log(-1): 순전파에서 NaN이 생기고 역전파로 전파돼.
0 / 0: 노름으로 정규화할 때 그 노름이 0이면 생겨.
0 * log(0): 0 × (-∞)의 결과가 NaN이 돼.
log(softmax(x))에서 소프트맥스 값 하나가 0에 가까워지면 역전파가 NaN이 될 수 있어. 수치적으로 안정적인 F.log_softmax를 사용해. 연습문제에서 직접 확인할 거야.
fp16 오버플로가 0×∞로 이어질 수 있어. bf16으로 바꾸거나 GradScaler를 추가해.

Code

set_detect_anomaly: Inf에서는 오류가 발생하지 않는 함정·python

import torch

# Turn on at the start of training when debugging — turn OFF for real runs
torch.autograd.set_detect_anomaly(True)

# Inf trap: anomaly detection WON'T catch this — sqrt(0) backward = inf, not nan
x = torch.tensor(0.0, requires_grad=True)
y = torch.sqrt(x)
y.backward()
print(x.grad, torch.isfinite(x.grad))   # tensor(inf), tensor(False) — no RuntimeError
# Use isfinite() to catch Inf; anomaly detection only fires on NaN.

set_detect_anomaly: NaN 역전파에서는 오류가 발생해·python

import torch

torch.autograd.set_detect_anomaly(True)

# Same sqrt function, different domain — sqrt(-1) IS caught
x = torch.tensor(-1.0, requires_grad=True)
y = torch.sqrt(x)            # forward: nan, propagates to backward
y.backward()
# RuntimeError: Function 'SqrtBackward0' returned nan values in its 0th output.
# The traceback points at the exact forward call that created the nan-producing op.

매개변수별 기울기 통계: 진단 순회·python

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 10))
loss = nn.functional.cross_entropy(model(torch.randn(8, 100)), torch.randint(0, 10, (8,)))
loss.backward()

print(f"{'name':25s}{'mean':>12s}{'std':>12s}{'max_abs':>12s}{'has_nan':>10s}")
for name, p in model.named_parameters():
    if p.grad is None:
        print(f"{name:25s}  NO GRADIENT")
        continue
    g = p.grad
    print(f"{name:25s}{g.mean():12.2e}{g.std():12.2e}{g.abs().max():12.2e}"
          f"{str(torch.isnan(g).any().item()):>10s}")

싼 NaN 보험: 모든 학습 반복문·python

import torch
import torch.nn as nn

model = nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for step, (x, y) in enumerate(loader):
    optimizer.zero_grad()
    loss = nn.functional.mse_loss(model(x), y)

    # Cheap. Worth it. isfinite() catches both Inf and NaN.
    if not torch.isfinite(loss):
        print(f"Step {step}: non-finite loss = {loss.item()}; halting.")
        # Save a snapshot of the offending batch for later analysis
        torch.save({'x': x, 'y': y, 'state': model.state_dict()},
                   f'nan_snapshot_step_{step}.pt')
        break

    loss.backward()
    optimizer.step()

Exercise

일부러 NaN을 만들어 봐. 로짓에 소프트맥스를 적용한 뒤 결과에 로그를 취하고, autograd.set_detect_anomaly(True)를 켜 역전파가 실패하게 해. 오류가 지목한 연산을 확인한 다음 log(softmax(x))를 F.log_softmax로 바꿔 문제를 해결해. 마지막으로 역전파가 정상적으로 끝나는지 확인해.

역전파가 잘못될 때 쓰는 도구

Inf와 NaN: 이상 탐지의 범위는 좁아

Code

External links

Exercise

Progress

댓글 0