[PYTORCH] 모델 학습 중 Loss NaN 발생 시 7가지 체크리스트와 즉시 해결 방법

모델 학습 중 Loss NaN 발생

 

딥러닝 엔지니어를 괴롭히는 "Not a Number" 현상, 원인 분석부터 수치적 안정성 확보까지

---

1. 서론: 왜 나의 모델은 NaN(Not a Number)을 뱉는가?

PyTorch를 이용해 야심 차게 모델 학습을 시작했는데, 수십 에폭(Epoch) 잘 돌아가던 중 갑자기 Loss가 NaN으로 변하는 순간은 모든 개발자에게 공포입니다. NaN은 한 번 발생하면 연쇄적으로 모든 가중치를 오염시키며 학습을 불가능하게 만듭니다. 이 현상은 단순한 버그가 아니라 수치적 불안정성(Numerical Instability)의 결과인 경우가 많습니다. 본 가이드는 실무에서 마주치는 NaN의 근본 원인을 해부하고, 이를 해결하기 위한 전략적 접근법을 제시합니다.

2. Loss NaN 발생의 주요 원인 비교 분석

NaN이 발생하는 지점은 데이터 입력, 연산 과정, 그리고 손실 함수 계산 단계로 나뉩니다. 각 원인별 특성과 차이를 표로 정리했습니다.

원인 분류	상세 원인	발생 시점 및 특징	해결 난이도
데이터 입력	Input 데이터 내 NaN/Inf 존재	학습 시작 직후 즉시 발생	낮음
경사도(Gradient)	Gradient Exploding (기울기 폭주)	학습 중간에 갑작스럽게 발생	중간
연산 안정성	0으로 나누기, log(0) 연산	특정 데이터 샘플 유입 시 발생	중간
하이퍼파라미터	과도하게 높은 Learning Rate	Loss가 요동치다 결국 NaN 수렴	낮음
정밀도 문제	Mixed Precision (FP16) Underflow	고성능 GPU 학습 시 간헐적 발생	높음

3. 실무자를 위한 NaN 해결 및 디버깅 Example 7가지

현업에서 즉시 적용할 수 있는 PyTorch 코드 기반의 해결책입니다.

Example 1: 입력 데이터 무결성 검사 (Sanity Check)

가장 기본적이지만 놓치기 쉬운 데이터 전처리 단계에서의 해결 방법입니다.

def check_data(dataloader):
    for batch_idx, (data, target) in enumerate(dataloader):
        if torch.isnan(data).any() or torch.isinf(data).any():
            print(f"NaN found in batch {batch_idx}")
            return False
    return True
        

Example 2: PyTorch 내장 디버깅 도구 활용 (anomaly_detection)

어느 연산에서 NaN이 시작되었는지 역추적하는 강력한 해결 방법입니다.

# 학습 루프 시작 전 설정
torch.autograd.set_detect_anomaly(True)

# 이후 연산 중 NaN 발생 시 역전파 단계에서 에러 메시지와 함께 중단됨
loss.backward() 
        

Example 3: Gradient Clipping을 이용한 폭주 방지

기울기가 너무 커져 가중치가 파괴되는 현상을 막는 표준적인 방법입니다.

loss.backward()
# L2 Norm 기준으로 1.0을 넘지 않도록 제한
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
        

Example 4: 수치적 안정성을 고려한 Log 연산 (Epsilon 추가)

log(0)이 발생하는 상황을 방지하기 위해 아주 작은 값(epsilon)을 더해주는 해결책입니다.

# 불안정한 방식
loss = -torch.log(prob)

# 안전한 방식 (eps=1e-7)
eps = 1e-7
loss = -torch.log(prob + eps)
        

Example 5: Softmax 대신 LogSoftmax 사용하기

지수 연산의 폭주를 막기 위해 PyTorch가 권장하는 함수 간의 차이를 활용한 방법입니다.

# 비권장: Softmax 연산 후 다시 Log를 취하는 것은 불안정함
# 권장: 내부적으로 최적화된 LogSoftmax와 NLLLoss 조합 사용
criterion = torch.nn.CrossEntropyLoss() # 이미 내부에서 안정적 연산 수행
        

Example 6: Mixed Precision 학습 시 GradScaler 적용

FP16 학습 시 발생하는 Underflow를 방지하여 NaN을 해결하는 방법입니다.

scaler = torch.cuda.amp.GradScaler()

for data, target in dataloader:
    with torch.cuda.amp.autocast():
        output = model(data)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
        

Example 7: 가중치 초기화(Weight Initialization) 전략 수정

초기 가중치가 너무 커서 발생하는 발산을 Xavier나 He 초기화로 해결하는 방법입니다.

def init_weights(m):
    if isinstance(m, torch.nn.Linear):
        torch.nn.init.kaiming_normal_(m.weight)
        m.bias.data.fill_(0.01)

model.apply(init_weights)
        

4. 결론 및 SEO 최적화 제언

Loss NaN은 단순한 오타가 아니라 하드웨어 정밀도와 수학적 한계가 부딪히는 지점에서 발생합니다. 가장 먼저 데이터의 정규화(Normalization) 상태를 확인하고, 그 다음으로 Learning Rate를 낮춰보십시오. 만약 여전히 해결되지 않는다면 PyTorch의 anomaly_detection을 활성화하여 범인(연산)을 찾아내는 과학적인 접근이 필요합니다.

5. 출처 및 참고 문헌

• PyTorch Documentation: "Autograd Mechanics - Anomaly Detection"
• NVIDIA Deep Learning Institute: "Numerical Stability in Neural Networks"
• Bengio et al., "Deep Learning", MIT Press (Chapter 7: Regularization for Deep Learning)