[PyTorch] torch.nn.KLDivLoss() 사용법과 예제

 

어쩌다보니 처음으로 KLDivergence를 loss로 사용하게 되었다. keras에서는 그냥 평범(?)하게 사용했으나 PyTorch에서는 우여곡절 끝에 제대로 loss가 떨어지는 것을 확인했다. 그래서 겸사겸사 기록해둔다.

KLDivLoss 선언하기 

처음엔 그냥 간단하게 선언을 했으나 reduction이라는 parameter가 mean으로 설정되어있는데 이렇게 설정되어 있으면 진짜 KLDivergence value를 뱉지 않는다고 한다. 그러므로 선언할 때부터 아래와 같이 선언해주어야 한다. 다음 버전에서는 이 부분을 개선하여 업데이트한다고 한다.

 

import torch

Loss = torch.nn.KLDivLoss(reduction='batchmean')

모델 정의하기

제대로 된 모델을 정의한 것은 아닌데 마지막 return을 보면, softmax를 쓰는 것이 아닌 log_softmax를 쓴 것을 확인할 수 있다. PyTorch에서 KLDivergence를 쓰려면 y_estimated에 log를 취한 값과 y_target을 loss에 넣어주어야 한다. 그래서 아래와 같이 log_softmax를 사용하면 된다.

 

class classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Conv1d(in_channels=1, out_channels=256,kernel_size=3)
        self.fc1 = nn.Linear(256,3)

        
    def forward(self, x):
    
        x = F.relu(self.layer1(x))
        x = F.relu(self.fc1(x))
        x = self.drop_out(x)
        x = self.fc2(x)
        
        return F.log_softmax(x, dim=1)

test할 때는 어떻게 해야하나?

마지막에 학습된 model을 이용해서 test를 해야하는데 이때 확률값을 얻고싶으나 log가 취해진 값이 나올 것이다. 그러므로 다시 exp()함수를 이용해서 log를 벗긴 값을 사용해야 한다.

 

(추가) 더 간단한 방법

log_softmax 쓰지않고 softmax쓴 이후에 loss넣기 전에 log 씌워주면 된다.

 

...
y = F.softmax(x,dim=1)
return y


criterion = torch.nn.KLDivLoss(reduction='batchmean')
loss = criterion(y.log(), y_real)