I'm experiencing an interesting and frustrating issue with the Dice loss used in image segmentation with Unet.
I have to segment images in two classes: background and region of interest. The region of interest is typically 4% of the pixels of the whole image. Images are about 1600x1600 pixels. I found the Dice loss working much better than Cross Entropy. However, if I use the standard Dice loss formula my Unet does not provide a correct output, i.e. all the pixels are predicted as background.
With standard Dice loss I mean:
where x_{c,i} is the probability predicted by Unet for pixel i and for channel c, and y_{c,i} is the corresponding ground-truth label. The modified version I use is:
Note the squared x at the denominator.
For some reason the latter one makes the net to produce a correct output, although the loss converges to ~0.5.
I do not understand why the latter works and the former doesn't. The latter works even if I use the power of three at the denominator.
Here below my implementation:
def make_one_hot(labels, classes):
one_hot = torch.FloatTensor(labels.size()[0], classes, labels.size()[2], labels.size()[3]).zero_().to(labels.device)
target = one_hot.scatter_(1, labels.data, 1)
return target
class DiceLoss(nn.Module):
def __init__(self,):
super(DiceLoss, self).__init__()
def forward(self, output, target):
target = make_one_hot(target.unsqueeze(dim=1), classes=output.size()[1])
output = F.softmax(output, dim=1)
numerator = (output * target).sum(dim=(2, 3))
denominator = output.pow(2).sum(dim=(2, 3)) + target.sum(dim=(2, 3))
iou = numerator / denominator
return 1 - iou.mean()


