ValueError: Expected target size (5, 50, 76, 80), got torch.Size([5, 2])

1

I am using a transfer learning from a pretrained model. When I tried to do this:

loss_f = nn.CrossEntropyLoss(weight=torch.tensor([1,1,1]))
output = model(volumes)
     
labels = torch.FloatTensor(labels).long()   #deleted .cuda()
labels = labels.cuda()

loss = loss_f(output, labels)

The error came up with:

ValueError: Expected target size (5, 50, 76, 80), got torch.Size([5, 2])

I printed the output and labels shape as follows:

output_shape = torch.Size([5, 2, 50 ,76, 80)]
labels_shape = torch.Size([5, 2])

It seems that with the default documentation from CrossEntropyLoss is (N, C, D1, D2, …, Dk), which k is the dimensions of the image. I am trying to do binary classification, so the classes are 2. I am facing the problem with the shape differences between the output and the labels, could anyone comment on how to solve the problem.

Thanks!!

2

The output shape indicates a multi-classification use case for volumetric data as it would correspond to logits in [batch_size, nb_classes, depth, height, width].
If you are working an a binary classification use case the output is expected to have the shape [batch_size, nb_classes2] while the target should have the shape [batch_size] and contain values in the range [0, nb_classes=2].
Currently your label tensor also seems to be one-hot encoded so use labels = torch.argmax(labels,. dim=1) if that’s the case.