Semantic segmentation loss function questions

vision
1

Hi there, I’m new to segmentation model.
I would like to use the deeplabv3_resnet50 model.
My image has shape (256, 256, 3) and my label has shape (256, 256). Each pixel in my label has a class value(0-4). And the batch size set in the DataLoader is 32.
The shape of my input batch is [32, 3, 256, 256] and the shape of corresponding target is [32, 256, 256]. I believe this is correct.
I was trying to use nn.BCEWithLogitsLoss().

  1. Is this the correct loss function for my case? Or should I use CrossEntropy instead?
  2. If this is the right one, the output of my model is [32, 5, 256, 256]. Each image prediction has the shape [5, 256, 256], does layer 0 means the unnomarlized probabilities of class 0? In order to make a [32, 256, 256] tensor to match the target to feed into the BCEWithLogitsLoss, do I need to transform the unnomarlized probabilities to classes?

Thank you everyone.