<p>Milletari et al. already explain this when they propose this in <a href="https://arxiv.org/pdf/1606.04797.pdf" rel="nofollow noreferrer">the paper of V-Net</a>. 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.</p>
<p><a href="https://i.stack.imgur.com/L9P4K.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/L9P4K.png" alt="enter image description here" /></a></p>
Dice loss working only when probs are squared at denominator
<p>I'm experiencing an interesting and frustrating issue with the Dice loss used in image segmentation with Unet.</p>
<p>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.</p>
<p>With standard Dice loss I mean:</p>
<p><a href="https://i.stack.imgur.com/HO2Ix.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/HO2Ix.png" alt="enter image description here" /></a></p>
<p>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:</p>
<p><a href="https://i.stack.imgur.com/srLBM.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/srLBM.png" alt="enter image description here" /></a></p>
<p>Note the squared x at the denominator.</p>
<p>For some reason the latter one makes the net to produce a correct output, although the loss converges to ~0.5.</p>
<p>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.</p>
<p>Here below my implementation:</p>
<pre><code>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()
</code></pre>