My final segmentation results are good even though my loss is very high?

Im using a Unet network for semantic segmentation. The loss reported by the network is very poor and so is the accuracy which is simple the intersection over union. Even after tens of thousands of iterations the loss does not seem to go below 0.91.

Yet, if I manually see the results the network seems to have segmented everything correctly and as expected.

What may be the reason for this, I know that a pytorch model returns logits and that is what is what is expected by the loss function or is my understanding not accurate?

I get the prediction outputs by applying a nn.Sigmoid on the final convolution layer output but during training this is not applied as, like I said, I thought the loss function expected logits without the Sigmoid.

Why is the loss not a representation of the final results?

Ive tried with nn.BCEWithLogitsLoss and DiceLoss

Thanks

After reading this: Calculating accuracy for a multi-label classification problem I managed to get the accuracy measure working by thresholding the network output. I can now see why it was not working before as my target was in binary format of 1 and 0 while the output from the nn.Sigmoid layer was a float value at every location. Thresholding now produces a more accurate IoU measure.

This now leaves the loss, with the above change my accuracy is improving as the network trains but the loss stays high. I am not passing the final convolution output from the network through a sigmoid layer and I am not thresholding it. The loss is given directly as from the model output:

output = model(input)
loss = loss_criterion(output, target)

thresholded_output = nn.Sigmoid(output) > 0.5
accuracy = accuracy_criterion(thresholded_output.to(dtype=torch.float), target)

Where target is the same binary tensor of 1 and 0 (dtype=Float32)

Should I be making any change before feeding the network output to loss criterion as well?

Thanks

nn.BCEWithLogitsLoss expects your model output to be logits, so you shouldn’t use a threshold on them.
On the other side, your custom DiceLoss might expect class predictions, but that depends on your implementation of it.

2 Likes