UNET with CrossEntropy Loss Function

Question

I was trying to train UNET with input size as [3,128,128] and the corresponding mask is [1,128,128] which contains classes directly(instead of pixels it will contain class numbers - 1,2). I am trying for a two-class problem hence my mask contains 1,2 as labels. Now I send my images to the model and the dimension of the predicted masks are [2,128,128]. Now to train a model I choose 16 as batch size. So, now I have input as [16,3,128,128] so the predicted dimension is [16,2,128,128]. But I have ground-truth masks as [16,1,128,128]. Now how can I apply Cross entropy loss in Pytorch? I have tried as follows and getting the following error. Could you please help? Thanks in advance.

lr = 0.1  # 0.1
criterion = nn.CrossEntropyLoss() #nn.L1Loss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, nesterov=True, weight_decay=0.0001)
is_train = True
is_pretrain = False
acc_best = 0
total_epoch = 30
if is_train is True:
    # Training
    for epoch in range(total_epoch):
        model.train()
        tims = time.time()
        for i, (images, labels) in enumerate(train_loader):
            images = Variable(images.permute(0,3,1,2).float().cuda())                                                                                                   
            labels = Variable(labels.type(torch.LongTensor).cuda())
            # Forward + Backward + Optimize
            optimizer.zero_grad()
            outputs = model(images)
            outputs = outputs.type(torch.float)
            print('predictedLabelsType:',outputs[0].type())
            print('ActualLabelsType:',labels[0].type())
            print('shape of predicted outputs:',outputs.shape)
            print('shape of groundtruth masks:',labels.shape)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

My output is as follows:

predictedLabelsType: torch.cuda.FloatTensor
ActualLabelsType: torch.cuda.LongTensor
shape of predicted outputs: torch.Size([16, 2, 128, 128])
shape of groundtruth masks: torch.Size([16, 1, 128, 128])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-8-b692a8d536a9> in <module>()
     52             print('shape of predicted outputs:',outputs.shape)
     53             print('shape of groundtruth masks:',labels.shape)
---> 54             loss = criterion(outputs, labels)
     55             loss.backward()
     56             optimizer.step()

3 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2385         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2386     elif dim == 4:
-> 2387         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2388     else:
   2389         # dim == 3 or dim > 4

RuntimeError: 1only batches of spatial targets supported (3D tensors) but got targets of size: : [16, 1, 128, 128]

Could you please suggest where is the mistake and how this CrossEntropyLoss() will work in Pytorch for image segmentation? What I am missing here? I have tried reshaping the target size as [16,128,128] which is leads to another error. Thanks a lot!

score 1 · Answer 1 · answered Apr 29 '21 at 18:52

The documentation specifies that if the input is shape (N, C, d1, d2) then the target must be shape (N, d1, d2). Instead, your targets are shape (N, 1, d1, d2) so you need to remove the unnecessary unitary dimension.

loss = criterion(output, labels.squeeze(1))

If you're getting another error from this change then there's another issue with your code, but this is the correct tensor shape for CrossEntropyLoss.

UNET with CrossEntropy Loss Function

1 Answers1