[Pytorch] softmax와 log_softmax (그리고 CrossEntropyLoss)

Pytorch로 MNIST 분류 예제 문제를 구현하다가, torch.nn.functional에 softmax, log_softmax 두 가지가 있다는 것을 발견했습니다.

2020/12/01 - [ML & DL/pytorch] - [Pytorch] MNIST Classification (2020/12/02 수정)

[Pytorch] MNIST Classification (2020/12/02 수정)

그리고 CrossEntropyLoss를 위 예제에서 사용을 했는데, 잘못 사용했다는 것을 발견했고, softmax와 log_softmax를 비교하고 CrossEntropyLoss에 대해서 정리해보도록 하겠습니다.

결론부터 말하자면 CrossEntropyLoss에는 log_softmax 연산이 내부에서 수행됩니다. 따라서 CrossEntropyLoss의 input으로는 model의 raw output으로 사용되어야 합니다.

 

 

공식 홈페이지는 softmax는 다음과 같이 설명하고 있습니다.

따라서 softmax 연산을 수행하면, elements들의 0에서 1사이의 확률을 갖게 되고, elements의 총합은 1이 됩니다. 또한, softmax에 의한 결과는 NLL Loss에 직접적으로 사용될 수 없다고 합니다.(제대로 동작하지 않음) 수치적으로 불안정해서 NaN의 결과를 얻을 수 있으므로, log_softmax와 NLL Loss를 결합하는 것을 추천합니다.

 

그리고 log_softmax는 다음과 같습니다.

softmax 함수에 log를 취한것과 동일합니다. 하지만, 두 연산을 따로 수행하는 것보다 더 빠르고 수치적으로 안정적이라고 합니다. 이 함수는 output과 gradient를 올바르게 구하기 위한 대안 공식(alternative formulation)을 사용합니다.

 

이 두 함수를 pytorch에서 직접 사용해서 비교해보도록 하죠.

[1,2,3,4,5]로 이루어진 벡터 텐서를 정의하고 softmax연산을 적용하면, 0~1사이의 확률을 갖는 벡터로 변환되고, 그 합은 1이라는 것을 볼 수 있습니다.

이번에는 (3, 5) 차원을 갖는 행렬 텐서를 정의하고 softmax를 적용했더니, 각 행 별로 총합이 1인 (3, 5)의 확률을 갖는 행렬로 변환되었습니다.

 

x에 log_softmax를 적용(z)해보고, y에는 log를 적용해보도록 하겠습니다.

두 결과값이 동일하다는 것을 볼 수 있습니다.

즉, log_softmax는 softmax에 log를 취한것과 동일하다는 것을 보여주고 있습니다.

 

이번에는 CrossEntropyLoss에 대해서 살펴보겠습니다.

공식 홈페이지에 길게 설명이 되어있는데, 핵심은 아래 공식입니다.

log_softmax와 nll_loss가 하나의 함수로 합쳐진 것이죠.

 

실제로 그런지 직접 계산해보도록 하겠습니다.

우선 label 역할의 y와 y에 one hot encoding 처리를 한 y_one_hot을 새로 정의하겠습니다.

class label이 3, 1, 0 이기 때문에 one hot encoding을 통해서 각 행별로 3, 1, 0 인덱스의 확률이 1이되고 나머지는 전부 0입니다.

 

그리고 아까전에 사용했던 x를 가지고 그대로 구해보겠습니다.

첫 번째는 log_softmax + nll_loss 입니다.

 

그리고, cross_entropy만을 사용한 loss입니다. cross_entropy는 내부에서 log_softmax 연산이 수행되기 때문에 x를 바로 input으로 사용합니다.

두 결과가 동일한 것을 볼 수 있습니다.

 

 

그리고 이전 MNIST 예제에서 Loss로 CrossEntropyLoss를 사용한다면, model 마지막 output에 log_softmax를 사용하면 안된다고 했는데, 엄밀히 말하자면 사용하지 말아야되는 것은 아닙니다.

 

이는 log_softmax(log_softmax(x)) = log_softmax(x) 이기 때문인데, 수학적인 증명은 하지 않도록 하고 실제 연산에서 동일한지 확인해보도록 하겠습니다.

보다시피 두 연산이 동일한 결과를 보여주고 있습니다.

 

따라서 CrossEntropyLoss를 사용하지만, 무조건 input으로 model raw output을 사용할 필요는 없는 것이죠. 

하지만 모델이 deep해질수록 조금이라도 모델을 줄이는 것이 유리해보이므로, 잘 알아보고 사용하는 것이 중요한 것 같습니다.

 

 

참고 : wikidocs.net/60572