Loss function returned Nan values 해결

 

직접 model architecture와 dataset, loss function을 만들고 학습시키려고 보니 loss function returned Nan values와 같은 에러가 발생했습니다.
아래와 같은 순서대로 어디서부터 문제가 발생했는지 확인하고, 해결할 수 있었습니다.

1. Learning rate 줄여보기

Learning rate이 너무 크면 모델이 loss function의 minimum value로 optimize되지 못하고 oscillate / overshoot 할 수 있습니다.
따라서 안정적으로 converge하도록 하려면 더 작은 learning rate으로 조정해볼 필요가 있습니다.
ex) 0.001 → 0.0001

 

2. Model architecture 또는 parameter 확인

모델 forward pass의 output 중 NaN 값이 발생하는 경우가 있는지 확인해봅니다.

for param in model.parameters():
    if torch.isnan(param).any():
        print("NaN detected in model parameters!")

 

3. Gradient Clipping 시도

gradient explosion을 방지하기 위해 clipping을 시도해볼 수 있습니다. 아래 코드와 같이 train 과정 중에 적용하면 됩니다.
하지만 근본적인 해결책은 아닌 것 같아서 어떤 상황에서 유용한지는 잘 모르겠네요!

max_norm 값은 gradient norm의 최대값을 지정하는 parameter입니다.
우선 1.0에서 5.0과 같은 작은 값에서 시작해서 먼저 시도해보고,
해당 값이 너무 작아서 learning process를 방해한다고 판단되면 10.0 또는 그 이상으로 늘려보면서 실험해보면 됩니다.

max_norm = 10

# in training phase
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
optimizer.step()

 

https://miro.medium.com/v2/resize:fit:1400/0*xUe5HglgQntK4aVJ.jpg

 

4. Gradient Anomaly Detect 시도

결국 제가 문제의 원인을 찾아낸 방법은 이것이었습니다.
아래 코드처럼 training process 시작 전에 torch.autograd.set_detect_anomaly(True)를 적용합니다.

torch.autograd.set_detect_anomaly(True)
for epoch in range(epochs):
    for imgs, targets in train_loader:
        # training process

그랬더니 다음과 같이 어느 부분에서 nan values가 생기는지 알려줍니다.
제 경우에는 Loss 계산 중에 Mse loss를 계산하는 부분이 문제였습니다.

MseLoss returned nan values 에러 화면

 

보통 MSELoss 계산 시 nan values가 return되는 이유는 크게 2가지인 것 같습니다.
zero division 같은 numerical instability, 또는 input data에 오류가 있는 것입니다.

두 경우 모두 디버깅 과정을 통해 원인을 파악하고 해결해야 합니다. 제 경우에는 input data를 처리하는 과정에서 오류가 있었습니다.
문제를 모두 해결하고 나니 잘 돌아가네요!

training 화면