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

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

 

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

 

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)   nn.BCELoss   사용

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 함수 입력값으로 넣으면   RuntimeError  가 발생한다.

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

 

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

그 뒤 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)   nn.BCEWithLogitsLoss   사용

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

 

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

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

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