Dice loss working only when probs are squared at denominator

Viewed 571

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:

enter image description here

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:

enter image description here

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()
1 Answers

Milletari et al. already explain this when they propose this in the paper of V-Net. They suggest that the ROI may only occupy a very small region of the whole scan, which is likely to be biased towards the background. Since you say your ROI is about 4% of the whole image, maybe you're facing a similar issue.

enter image description here

Related