What is the loss function for multilabel segmentation

[talhaanwarch]

I am confused what loss function should I used. My image size and label size is (512,512,128). The maximum value in my label is 4 and the minimum is 0. So I have 5 classes.
Here is my UNET

UNet(
                      dimensions=3,
                      in_channels=1,
                      out_channels=5,
                      channels=(16, 32, 64, 128, 256),
                      strides=(2, 2, 2, 2),
                      num_res_units=2,
                      norm=Norm.BATCH,
                  )

Here is loss function, I am using
DiceLoss(to_onehot_y=True, sigmoid=True, squared_pred=True)
But my output produced from the model is not between 0 and 4. Its around -100 to 100
I am printing shape, max value and min value below for both input to model and output from model.

train  input torch.Size([1, 1, 512, 512, 128]) tensor(4., device='cuda:0') tensor(0., device='cuda:0')

train  output torch.Size([1, 5, 512, 512, 128]) tensor(22.9688, device='cuda:0', dtype=torch.float16, grad_fn=<MaxBackward1>) tensor(-7.4922, device='cuda:0', dtype=torch.float16, grad_fn=<MinBackward1>)

[Nic-Ma]

Hi @talhaanwarch ,

Thanks for your interest and experiments here.
I think that's the expected output? Because you don't have sigmoid and argmax in UNet, DiceLoss has sigmoid and to_onehot_y args to handle the raw output of UNet as you set.

Thanks.