I am trying to do a binary segmentation task on the coco dataset using RestNet18 as encoder and DeepLabV3 as the decoder. Looking around i found that BCEWithLogitsLoss is the recommended criterion but i cannot get my training to work. Can you please help me proceed in the right direction ?
I am trying the following in the train method :
outputs = model(inputs)
loss = criterion(torch.nn.Sigmoid(outputs), labels)
The input to my network is torch.Size([32, 3, 256, 256])
The Output is torch.Size([32, 2, 256, 256])
and the labels are torch.Size([32, 1, 256, 256])
I cannot get the sigmoid to work, earlier i tried the CrossEntropyLoss with 2 classes but i doubt if it even worked. I am mostly stressed about the shapes.
Will be grateful for any help, thanks !