[PYTORCH] nn.CrossEntropyLoss 사용 시 Softmax 중복 적용을 피하는 2가지 해결 방법과 성능 차이

nn.CrossEntropyLoss 사용 시 Softmax 중복 적용

 

 

PyTorch 프레임워크를 사용하여 딥러닝 모델을 설계할 때, 입문자와 실무자 모두가 가장 빈번하게 실수하는 지점 중 하나가 바로 다중 클래스 분류(Multi-class Classification) 모델의 출력층 설계입니다. 특히 nn.CrossEntropyLoss를 사용할 때 Softmax 함수를 명시적으로 적용해야 하는지에 대한 의문은 모델의 수렴 속도와 수치적 안정성에 직결되는 매우 중요한 문제입니다. 결론부터 말씀드리면, PyTorch의 nn.CrossEntropyLoss는 내부적으로 LogSoftmax와 NLLLoss(Negative Log Likelihood Loss)를 결합하여 처리하므로, 모델의 마지막 레이어에 별도의 Softmax를 적용해서는 안 됩니다. 본 가이드에서는 이 설계 원리와 실무에서 발생할 수 있는 7가지 구체적인 구현 예시를 통해 완벽한 모델 최적화 방법을 제시합니다.

---

1. nn.CrossEntropyLoss의 내부 메커니즘과 동작 원리

딥러닝에서 손실 함수를 계산할 때 수치적 안정성(Numerical Stability)은 매우 중요합니다. 확률 값으로 변환된 $Softmax$ 출력값은 0과 1 사이의 값을 가지는데, 이 값이 매우 작아질 경우 로그 연산 과정에서 -inf로 발산하는 Underflow 문제가 발생할 수 있습니다. PyTorch 개발팀은 이를 방지하기 위해 nn.CrossEntropyLoss 내부에 LogSoftmax 연산을 통합했습니다. 따라서 사용자는 모델의 최종 출력값으로 가중치 합산 결과인 Logits(로짓)을 그대로 전달해야 합니다.

수학적 설계 차이 요약

모델의 최종 출력을 $z$라고 할 때, 각 방식의 차이는 다음과 같습니다.

구분	입력 형태 (Input)	내부 연산 과정	최종 출력 형태
nn.CrossEntropyLoss	Raw Logits (Unnormalized)	LogSoftmax + NLLLoss	Scalar Loss
nn.NLLLoss	Log-probabilities	Negative Log Likelihood	Scalar Loss
nn.BCELoss	Probabilities (0~1)	Binary Cross Entropy	Scalar Loss

---

2. 실무 적용을 위한 7가지 핵심 Example 코드

실제 프로젝트에서 발생할 수 있는 다양한 케이스를 고려하여, nn.CrossEntropyLoss를 올바르게 사용하는 코드 패턴을 정리했습니다.

Example 1: 표준적인 다중 클래스 분류 모델 구조

import torch.nn as nn

class StandardClassifier(nn.Module):
    def __init__(self):
        super(StandardClassifier, self).__init__()
        self.fc = nn.Linear(512, 10) # 10개의 클래스

    def forward(self, x):
        # 마지막에 Softmax를 사용하지 않고 Logits를 반환함
        return self.fc(x)

model = StandardClassifier()
criterion = nn.CrossEntropyLoss()

Example 2: 추론(Inference) 단계에서의 Softmax 적용

학습 시에는 로짓을 사용하지만, 사용자에게 확률을 보여줘야 하는 서비스 배포 시에는 수동으로 적용합니다.

model.eval()
with torch.no_grad():
    logits = model(input_data)
    # 확률값을 얻기 위해 추론 시에만 Softmax 적용
    probabilities = torch.softmax(logits, dim=1)
    prediction = torch.argmax(probabilities, dim=1)

Example 3: nn.Sequential을 활용한 간결한 구성

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10) # Softmax 레이어를 추가하지 않음
)

Example 4: Label Smoothing 적용 방법

PyTorch 1.10 이상에서는 nn.CrossEntropyLoss 내에서 직접 라벨 스무딩을 지원합니다.

# 과적합 방지를 위한 0.1 비율의 Label Smoothing
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

Example 5: 클래스 불균형 해결을 위한 Weight 부여

# 클래스별 데이터 양이 다를 경우 가중치 설정
weights = torch.tensor([1.0, 2.0, 0.5, ...])
criterion = nn.CrossEntropyLoss(weight=weights)

Example 6: LogSoftmax + NLLLoss 수동 분리 (커스텀 필요 시)

특정 레이어의 로그 확률값이 중간 과정에서 필요할 때 사용하는 방식입니다.

class CustomModel(nn.Module):
    def forward(self, x):
        logits = self.fc(x)
        return nn.functional.log_softmax(logits, dim=1)

model = CustomModel()
criterion = nn.NLLLoss() # CrossEntropy 대신 NLLLoss 사용

Example 7: Ignore Index를 활용한 패딩 처리

NLP 작업 등에서 특정 인덱스(예: Padding)를 손실 계산에서 제외할 때 유용합니다.

# 인덱스 0번은 손실 계산에서 무시
criterion = nn.CrossEntropyLoss(ignore_index=0)

---

3. Softmax를 중복 적용했을 때 발생하는 문제점

만약 모델의 마지막에 nn.Softmax()를 넣고 손실 함수로 nn.CrossEntropyLoss를 사용하면 어떤 일이 벌어질까요?

1. Gradient Vanishing (기기값 소실): Softmax 출력값은 0~1 사이로 압축됩니다. 여기에 다시 내부적으로 LogSoftmax를 취하면 값이 매우 작은 음수가 되어 기울기가 소실되고 학습 속도가 현저히 느려집니다.
2. 수치적 불안정성: 지수 함수의 중복 계산으로 인해 Floating Point 오차가 누적될 수 있습니다.
3. 성능 저하: 동일한 에포크를 학습하더라도 정확도(Accuracy)가 로짓을 직접 사용했을 때보다 낮게 나타나는 경향이 있습니다.

---

4. 결론 및 요약

PyTorch 실무에서 nn.CrossEntropyLoss는 "Logits to Loss" 공식만 기억하면 됩니다. 모델 코드 내부에 nn.Softmax가 보인다면, 그것이 학습을 위한 것인지 아니면 단순 출력을 위한 것인지 반드시 점검하십시오. 최신 딥러닝 아키텍처는 대부분 수치 안정성을 위해 로짓 단에서 손실을 계산하도록 설계되어 있습니다.

작업 단계	권장 사항 (Solution)
모델 학습 (Training)	마지막 레이어에서 활성화 함수 없이 Linear 출력(Logits)만 반환
손실 함수 정의	nn.CrossEntropyLoss() 사용
추론 및 배포 (Inference)	torch.softmax() 또는 argmax()를 명시적으로 적용

 

출처 및 참고문헌:
1. PyTorch Official Documentation: nn.CrossEntropyLoss (v2.2)
2. Deep Learning with PyTorch (Eli Stevens et al., Manning Publications)
3. Neural Networks and Deep Learning (Michael Nielsen)