[FIXED] What is the appropriate way to use BCE loss with ResNet outputs?

Issue

The title explains the overall problem, but for some elaboration:

I'm using torchvision.models.resnet18() to run an anomaly detection scheme. I initialize the model by doing:

net = torchvision.models.resnet18(num_classes=2)

since in my particular setting 0 equals normal samples and 1 equals anomalous samples.

The output from my model is of shape (16, 2) (batch size is 16) and labels are of size (16, 1). This gives me the error that the two input tensors are of inappropriate shape.

In order to solve this, I tried something like:

>>> new_output = torch.argmax(output, dim=1)

Which gives me the appropriate shape, but running loss = nn.BCELoss(new_output, labels) gives me the error:

RuntimeError: bool value of Tensor with more than one value is ambiguous

What is the appropriate way for me to approach this issue? Thanks.

Edit

I've also tried using nn.CrossEntropyLoss as well, but am getting the same RuntimeError.

More specifically, I tried nn.CrossEntropyLoss(output, label) and nn.CrossEntropyLoss(output, label.flatten()).

Solution

If you want to use BCELoss, the output shape should be (16, 1) instead of (16, 2) even though you have two classes. You may consider reading this excellent writing to under binary cross-entropy loss.

Since you are getting output from resnet18 with shape (16, 2), you should rather use CrossEntropyLoss where you can give (16, 2) output and label of shape (16).

You should use CrossEntropyLoss as follows.

loss_crit = nn.CrossEntropyLoss()
loss = loss_crit(output, label) 

where output = (16, 2) and label = (16). Please note, the label should contain either 0 or 1.

Please see the example provided (copied below) in the official documentation.

>>> loss = nn.CrossEntropyLoss()
>>> input = torch.randn(3, 5, requires_grad=True)
>>> target = torch.empty(3, dtype=torch.long).random_(5)
>>> output = loss(input, target)
>>> output.backward()

Answered By - Wasi Ahmad