[pytorch] How to use nn.CrossEntropyLoss() 사용법

 

 

 

아래 코드는 pytorch에서 loss function으로 CrossEntropy를 사용하는 예이다.

cls_loss = nn.CrossEntropyLoss()


test_pred_y = torch.Tensor([[2,0.1],[0,1]]) # 실제 사용에선 softmax에 의해 각 행의 합이 1이 될 것이다.
test_true_y1 = torch.Tensor([1,0]).long() # 1은 true값이 1번째(클래스)라는 것을 의미
test_true_y2 = torch.Tensor([0,1]).long()

print(test_pred_y)
print(test_true_y1)
print(test_true_y2)

print(cls_loss(test_pred_y, test_true_y1))
print(cls_loss(test_pred_y, test_true_y2))

 

실행시 결과 값이다.

# results
tensor([[2.0000, 0.1000],
        [0.0000, 1.0000]])
tensor([1, 0])
tensor([0, 1])
tensor(1.6763)
tensor(0.2263)

 

---

 

설명

cls_loss는 2개의 tensor 인자를 받는다.

하나는 예측된 y값, 다른 하나는 true y값이며, true_y의 형태에 주의해야한다.

 

 - pred y의 경우 softmax로 처리된 총합 1인 확률정보가 담겨있다.

    (각 값들의 총 합은 1이여야하지만, 1이 아닌경우도 유연하게 대처하는걸 확인 할 수 있다.)

 - true y의 경우 정수형으로 클래스(위치) 값이 담긴 텐서이다.

 

예를들어 아래  tensor의 경우,

# processed softmax
tensor([ [0.7, 0.2, 0.1]
         [0.3, 0.2, 0.5] 
       ])

true Y는 다음과 같을 것이다.

tensor([ [0, 2] ])