[pytorch] nn.BCEWithLogitsLoss VS nn.BCELoss 차이

[ic]nn.BCEWithLogitsLoss[/ic] 와 [ic]nn.BCELoss[/ic]는 둘 다 [ic]이진 분류(binary classification)[/ic]에 사용되는 로스함수다.

 

이번 포스팅에서는 두함수의 차이를 알아볼 거다.

 

1) 예제 코드 준비

(1) 모델 생성

import torch
import torch.nn as nn

# Define a simple model that takes in a single input and produces a single output
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)

# Create an instance of the model
model = Model()

 

(2) dummy 데이터 생성

# Dummy data
input_data = torch.tensor([[1.0], [2.0], [3.0]])
labels = torch.tensor([[0.0], [0.0], [1.0]])

 

2) [ic]nn.BCELoss[/ic] 사용

loss_fn = nn.BCELoss()
output = model(input_data)
loss = loss_fn(output, labels)
print(loss)  # logit을 activation function(sgimoid)를 감싸 0과 1사이로 변환한뒤 loss 함수에 넣어줘야함.

#output
'''
RuntimeError: all elements of input should be between 0 and 1
'''

 model에서 나온 output 값을 loss 함수 입력값으로 넣으면 [ic]RuntimeError[/ic]가 발생한다.

왜냐하면 [ic]nn.BCELoss[/ic]는 0과 1 사이 값만 받기 때문이다.

 

output을 activation function([ic]sigmoid[/ic] )을 통해 [ic]0과 1사이[/ic] 값으로 변환해줘야 한다.

그 뒤 loss함수 입력값으로 넣어주면 된다.

loss_fn = nn.BCELoss()
output = model(input_data)
output = torch.sigmoid(output)
loss = loss_fn(output, labels) 
print(loss)  

#output
'''
tensor(0.6418, grad_fn=<BinaryCrossEntropyBackward0>)
'''

 

3) [ic]nn.BCEWithLogitsLoss[/ic] 사용

[ic]nn.BCEWithLogitsLoss[/ic] 는 [ic]sigmoid[/ic] 함수가 포함되어 있는 [ic]nn.BCELoss[/ic]라고 이해하면 쉽다.

 

코드를 보면 [ic]sigmoid[/ic]  없이도 정상 작동된다.

loss_fn = nn.BCEWithLogitsLoss()
output = model(input_data)
loss = loss_fn(output, labels)
print(loss)

#output
'''
tensor(0.6418, grad_fn=<BinaryCrossEntropyBackward0>)
'''