Dice Loss with custom penalities

NearsightedCV: <p>Hi all,</p>
<p>I am wading through this CV problem and I am getting better results<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/1/31fbf0e29b170afeded6be4afd6f5b0a31e22b41.png" data-download-href="https://discuss.pytorch.org/uploads/default/31fbf0e29b170afeded6be4afd6f5b0a31e22b41" title="Screenshot 2022-03-13 204825"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31fbf0e29b170afeded6be4afd6f5b0a31e22b41_2_690x342.png" alt="Screenshot 2022-03-13 204825" data-base62-sha1="78baYDqvIb0GkhcDoT9Dbm4ZsFb" width="690" height="342" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31fbf0e29b170afeded6be4afd6f5b0a31e22b41_2_690x342.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31fbf0e29b170afeded6be4afd6f5b0a31e22b41_2_1035x513.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31fbf0e29b170afeded6be4afd6f5b0a31e22b41_2_1380x684.png 2x" data-dominant-color="75918C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2022-03-13 204825</span><span class="informations">1411×700 28.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The challenge is my images are imbalanced with background and one other class dominant. Cross Entropy was a wash but Dice Loss was showing some improvement in getting the less prevalent class but I think I need an added penalty on getting the less prevalent class wrong. They will always be less prevalent so I would just need to get the weights right.</p>
<p>I have been using the very handy segmentation models package but they do not provide a way of adding custom weights. So I want to modify the Dice Loss function with this upgrade.</p>
<p>I am confused where I should stick it in the code</p>
<pre><code class="lang-auto">from typing import Optional, List

import torch
import torch.nn.functional as F
from torch.nn.modules.loss import _Loss
from ._functional import soft_dice_score, to_tensor
from .constants import BINARY_MODE, MULTICLASS_MODE, MULTILABEL_MODE

__all__ = ["DiceLoss"]


[docs]class DiceLoss(_Loss):
    def __init__(
        self,
        mode: str,
        classes: Optional[List[int]] = None,
        log_loss: bool = False,
        from_logits: bool = True,
        smooth: float = 0.0,
        ignore_index: Optional[int] = None,
        eps: float = 1e-7,
    ):
        """Dice loss for image segmentation task.
        It supports binary, multiclass and multilabel cases

        Args:
            mode: Loss mode 'binary', 'multiclass' or 'multilabel'
            classes:  List of classes that contribute in loss computation. By default, all channels are included.
            log_loss: If True, loss computed as `- log(dice_coeff)`, otherwise `1 - dice_coeff`
            from_logits: If True, assumes input is raw logits
            smooth: Smoothness constant for dice coefficient (a)
            ignore_index: Label that indicates ignored pixels (does not contribute to loss)
            eps: A small epsilon for numerical stability to avoid zero division error
                (denominator will be always greater or equal to eps)

        Shape
             - **y_pred** - torch.Tensor of shape (N, C, H, W)
             - **y_true** - torch.Tensor of shape (N, H, W) or (N, C, H, W)

        Reference
            https://github.com/BloodAxe/pytorch-toolbelt
        """
        assert mode in {BINARY_MODE, MULTILABEL_MODE, MULTICLASS_MODE}
        super(DiceLoss, self).__init__()
        self.mode = mode
        if classes is not None:
            assert mode != BINARY_MODE, "Masking classes is not supported with mode=binary"
            classes = to_tensor(classes, dtype=torch.long)

        self.classes = classes
        self.from_logits = from_logits
        self.smooth = smooth
        self.eps = eps
        self.log_loss = log_loss
        self.ignore_index = ignore_index

    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -&gt; torch.Tensor:

        assert y_true.size(0) == y_pred.size(0)

        if self.from_logits:
            # Apply activations to get [0..1] class probabilities
            # Using Log-Exp as this gives more numerically stable result and does not cause vanishing gradient on
            # extreme values 0 and 1
            if self.mode == MULTICLASS_MODE:
                y_pred = y_pred.log_softmax(dim=1).exp()
            else:
                y_pred = F.logsigmoid(y_pred).exp()

        bs = y_true.size(0)
        num_classes = y_pred.size(1)
        dims = (0, 2)

        if self.mode == BINARY_MODE:
            y_true = y_true.view(bs, 1, -1)
            y_pred = y_pred.view(bs, 1, -1)

            if self.ignore_index is not None:
                mask = y_true != self.ignore_index
                y_pred = y_pred * mask
                y_true = y_true * mask

        if self.mode == MULTICLASS_MODE:
            y_true = y_true.view(bs, -1)
            y_pred = y_pred.view(bs, num_classes, -1)

            if self.ignore_index is not None:
                mask = y_true != self.ignore_index
                y_pred = y_pred * mask.unsqueeze(1)

                y_true = F.one_hot((y_true * mask).to(torch.long), num_classes)  # N,H*W -&gt; N,H*W, C
                y_true = y_true.permute(0, 2, 1) * mask.unsqueeze(1)  # H, C, H*W
            else:
                y_true = F.one_hot(y_true, num_classes)  # N,H*W -&gt; N,H*W, C
                y_true = y_true.permute(0, 2, 1)  # H, C, H*W

        if self.mode == MULTILABEL_MODE:
            y_true = y_true.view(bs, num_classes, -1)
            y_pred = y_pred.view(bs, num_classes, -1)

            if self.ignore_index is not None:
                mask = y_true != self.ignore_index
                y_pred = y_pred * mask
                y_true = y_true * mask

        scores = self.compute_score(y_pred, y_true.type_as(y_pred), smooth=self.smooth, eps=self.eps, dims=dims)

        if self.log_loss:
            loss = -torch.log(scores.clamp_min(self.eps))
        else:
            loss = 1.0 - scores

        # Dice loss is undefined for non-empty classes
        # So we zero contribution of channel that does not have true pixels
        # NOTE: A better workaround would be to use loss term `mean(y_pred)`
        # for this case, however it will be a modified jaccard loss

        mask = y_true.sum(dims) &gt; 0
        loss *= mask.to(loss.dtype)

        if self.classes is not None:
            loss = loss[self.classes]

        return self.aggregate_loss(loss)

    def aggregate_loss(self, loss):
        return loss.mean()

    def compute_score(self, output, target, smooth=0.0, eps=1e-7, dims=None) -&gt; torch.Tensor:
        return soft_dice_score(output, target, smooth, eps, dims)

</code></pre>
<p>Where can I add my edit? I see a reference to a</p>
<p><code> from ._functional import soft_dice_score</code><br>
I cannot find this any where. Ideally I would like to avoid writing something from scratch.  I suspect it would be where the DICE score is actually calculated but I cannot find it.</p>

KFrank: <p>Hi Matthew!</p>
<p>Please let me offer some suggestions without answering your specific<br>
question.</p>
<aside class="quote no-group quote-modified" data-username="NearsightedCV" data-post="1" data-topic="146334" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/n/ac8455/48.png" class="avatar"> NearsightedCV:</div>
<blockquote>
<p>The challenge is my images are imbalanced with background and one other class dominant. Cross Entropy was a wash but Dice Loss was shoing some improvement</p>
</blockquote>
</aside>
<p>I do realize that Dice Loss is often recommended for segmentation<br>
problems with unbalanced data, but let me suggest instead using<br>
<code>CrossEntropyLoss</code> with class weights (<code>CrossEntropyLoss</code>’s <code>weight</code><br>
constructor argument).</p>
<p>My advice is to start with (weighted) <code>CrossEntropyLoss</code>, and if<br>
that doesn’t seem to be doing well enough, try <em>adding</em> Dice Loss<br>
to <code>CrossEntropyLoss</code> as a further contribution to the total loss.<br>
My view is that doing so is likely to work better than using Dice Loss<br>
in isolation (and that weighted <code>CrossEntropyLoss</code> is likely to work<br>
well without adding Dice Loss).</p>
<aside class="quote no-group">
<blockquote>
<p>I have been using the very handy segmentation models package but they do not provide a way of adding custom weights. So I want to modify the Dice Loss function with this upgrade.</p>
</blockquote>
</aside>
<p>My only comment here is that if you have access to the code at a level<br>
where you could add custom weights to a Dice Loss, you should also be<br>
able to swap in <code>CrossEntropyLoss</code> (with its built-in <code>weight</code> argument).</p>
<p>Best.</p>
<p>K. Frank</p>

NearsightedCV: <p>Thanks for the tip! I will start with the custom weights and may add Dice later. I will keep you posted!!!</p>
<p><a class="mention" href="/u/kfrank">@KFrank</a></p>
<p>real quick question if I wanted to have both how would I do that?</p>

mMagmer: <p>Hi <a class="mention" href="/u/nearsightedcv">@NearsightedCV</a>,<br>
from what I know, dice loss for multi class is the average of dice loss for each class.<br>
So it is balancing data in a way.<br>
But if you want, I think you can change  how to average them.</p>
<aside class="quote no-group" data-username="NearsightedCV" data-post="1" data-topic="146334">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/n/ac8455/48.png" class="avatar"> NearsightedCV:</div>
<blockquote>
<pre><code class="lang-auto">def aggregate_loss(self, loss):
        return loss.mean()
</code></pre>
</blockquote>
</aside>
<p>Var <code>loss</code>  should be a vector with shape <code>#Classes</code>. You can multiply it with weight vector.<br>
But I also think you can get a good enough result with weighted cross entropy.</p>

NearsightedCV: <p><a class="mention" href="/u/mmagmer">@mMagmer</a> I was leaning towards Dice because just switching to it seemed to give a better result without weights. I started playing around with it, and it does seem to help.</p>
<p><a class="mention" href="/u/mmagmer">@mMagmer</a> do you know what the best way to optimize the weighting would be?</p>
<p>One other thing that is confusing me, using pretrained encoders vs starting from scratch. Do I still only have 10-20 epochs to make a model before overfitting risk takes place? I am getting training scores pretty low by 25 epochs but my validation scoring is worse.</p>
<p>This is</p>

mMagmer: <aside class="quote no-group" data-username="NearsightedCV" data-post="5" data-topic="146334">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/n/ac8455/48.png" class="avatar"> NearsightedCV:</div>
<blockquote>
<p>do you know what the best way to optimize the weighting would be</p>
</blockquote>
</aside>
<p>for cross entropy  i think it’s 1/<span class="hashtag-raw">#each_class_pixel</span><br>
for dice maybe you can find somting in <a href="https://arxiv.org/abs/1707.03237" class="inline-onebox" rel="noopener nofollow ugc">[1707.03237] Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations</a> .</p>
<blockquote>
<p>when choosing the GDLv weighting, the contribution of each label is<br>
corrected by the inverse of its volume, thus reducing the well known correlation<br>
between region size and Dice score.<br>
i really don’t know.</p>
</blockquote>
<p>when your training error is still high you’re not overfitting.</p>