(Weighted) Dice Loss only predicts 3 out of 9 classes

username: Tony_W
cooked: <p>Hello everyone, I’m kinda new to ML and CV and I’ve been training a semantic segmentation model for my master thesis. My model stagnates after 20ish epochs which it does not with CrossEntropyLoss. Also when testing out my model it only ever predicts the first 3 out of 9 classes. 2 of those classes are predominate in my dataset while one is actually relatively seldom.</p>
<p>This is my code for weighted dice loss:</p>
<pre><code class="lang-auto">class ClassAverageDiceLoss(nn.Module):

    def __init__(self, num_classes, softmax_dim=None):
        super().__init__()
        self.num_classes=torch.tensor(num_classes)
        self.softmax_dim=softmax_dim

    def forward(self, logits, targets, reduction='mean', smooth=1e-5):
        probabilities=logits
        if self.softmax_dim is not None:
            probabilities = nn.Softmax(dim=self.softmax_dim)(logits)
        #end if
        
        targets_one_hot=torch.nn.functional.one_hot(targets, num_classes=self.num_classes)

        # Convert from NHWC to NCHW
        targets_one_hot = targets_one_hot.permute(0, 3, 1, 2)

        dice_coeff_individual_array=torch.zeros(self.num_classes, dtype=torch.float32).to(DEVICE)

        #calcuate class specific dice loss/score
        for n_class in range(self.num_classes):

            targets_one_hot_class_specific=targets_one_hot[:,n_class,:,:]
            probabilities_class_specific=probabilities[:,n_class,:,:]

            intersection_class_specific = (targets_one_hot_class_specific * probabilities_class_specific).sum()

            mod_a = intersection_class_specific.sum()
                # .numel() and -log() in combination, due to the loss function always trying to minimize something
            mod_b = targets_one_hot_class_specific.numel()

            dice_coeff_individual = 2. * intersection_class_specific / (mod_a + mod_b + smooth)

            dice_coeff_individual_array[n_class] = dice_coeff_individual
        
        dice_coeff_average=dice_coeff_individual_array.sum()/(self.num_classes)

        dice_loss_average=-dice_coeff_average.log()

        return dice_loss_average
</code></pre>
<p>I thought that dice loss was good at handling imbalanced datasets which is why I was trying that out. If somebody could check my code for the weighted dice loss, I would be really thankful!</p>
<p>Here a graph of my training vs validation loss</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/3/235282722e6e7ef889404f08d05ee1714e8508c1.png" data-download-href="https://discuss.pytorch.org/uploads/default/235282722e6e7ef889404f08d05ee1714e8508c1" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/235282722e6e7ef889404f08d05ee1714e8508c1_2_517x337.png" alt="image" data-base62-sha1="52tudshsSChT6Wbor1fz5htr3tn" width="517" height="337" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/235282722e6e7ef889404f08d05ee1714e8508c1_2_517x337.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/235282722e6e7ef889404f08d05ee1714e8508c1_2_775x505.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/3/235282722e6e7ef889404f08d05ee1714e8508c1.png 2x" data-dominant-color="F9F8F8"><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">image</span><span class="informations">837×547 31 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>

username: KFrank
cooked: <p>Hi Tony!</p>
<aside class="quote no-group quote-modified" data-username="Tony_W" data-post="1" data-topic="202228" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tony_w/48/65459_2.png" class="avatar"> Tony_W:</div>
<blockquote>
<p>My model stagnates after 20ish epochs which it does not with CrossEntropyLoss. Also when testing out my model it only ever predicts the first 3 out of 9 classes. 2 of those classes are predominate in my dataset while one is actually relatively seldom.<br>
…<br>
I thought that dice loss was good at handling imbalanced datasets</p>
</blockquote>
</aside>
<p>(Note, I haven’t looked at your Dice-loss code.)</p>
<p>There is a legitimate argument that Dice loss should help with imbalanced datasets.</p>
<p>But <code>CrossEntropyLoss</code> is a very good loss for training classification models (and<br>
semantic segmentation is a kind of classification in that you are classifying pixels).<br>
My intuition is that the benefits of <code>CrossEntropyLoss</code> – once you incorporate class<br>
weights to address the issue of class imbalance – outweigh the benefits of Dice loss.</p>
<p>What is your goal here?  Why did you switch to Dice loss from <code>CrossEntropyLoss</code><br>
(when <code>CrossEntropyLoss</code> seemed to be working better)?  Is your project to compare<br>
how well the two loss functions work or is your primary goal to train an effective<br>
semantic-segmentation model?</p>
<p>My recommendation would be to use <code>CrossEntropyLoss</code> with class weights.  If that<br>
doesn’t work well enough and you think that Dice loss might help, I would <em>add</em> Dice<br>
loss (with a tunable hyperparameter weight) to <code>CrossEntropyLoss</code>.  That way you<br>
get the generally good convergence properties of <code>CrossEntropyLoss</code> together with<br>
any potential class-imbalance benefit from Dice loss.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Tony_W
cooked: <p>Hi K. Frank, thanks for the quick response. I’m trying out dice loss, because my final evaluation metric is dice score and I thought that directly training with dice loss would have a positive effect on the entire procedure.</p>
<p><em>What is your goal here? Why did you switch to Dice loss from <code>CrossEntropyLoss</code></em><br>
<em>(when <code>CrossEntropyLoss</code> seemed to be working better)?</em><br>
 → Currently, I can’t say if CrossEntropyLoss works better, because I have the feeling that my DiceLoss isn’t really working as intended.</p>
<p><em>Is your project to compare how well the two loss functions work or is your primary goal to train an effective semantic-segmentation model?</em><br>
 → I just want to train an effective segmentation model for my application. But since I’m also still learning, it would be interesting for me to utilize different loss functions and see how they compare.</p>
<p><em>I would <em>add</em> Diceloss (with a tunable hyperparameter weight) to <code>CrossEntropyLoss</code>.</em><br>
 → Okay, so a combination of Diceloss and CrossEntropyLoss? Would you just sum up the loss? And if so, would you still apply class weights to CrossEntropyLoss?</p>
<p>Thanks in advance!</p>
<p>Best regards, Tony</p>

username: KFrank
cooked: <p>Hi Tony!</p>
<aside class="quote no-group" data-username="Tony_W" data-post="3" data-topic="202228" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tony_w/48/65459_2.png" class="avatar"> Tony_W:</div>
<blockquote>
<p>my final evaluation metric is dice score and I thought that directly training with dice loss would have a positive effect on the entire procedure.</p>
</blockquote>
</aside>
<p>This is not an unreasonable way of looking at things and it often makes sense.<br>
However, it can also make sense to use a loss function for training that, while<br>
related to, is not the same as the final metric you use to decide how good your<br>
model is.  It’s perfectly possible for your training loss to have, so to speak, better<br>
training properties than your evaluation metric does.</p>
<p>As a simple example, if your final metric isn’t differentiable, you can’t train with<br>
it (using gradient-descent optimizers).</p>
<p>Consider ordinary classification (as distinct from semantic-segmentation pixel<br>
classification).  A common metric is the accuracy – the fraction of samples that<br>
are correctly classified.  As this involves discrete counting, it’s not differentiable.<br>
But it is also possible to define a “probabilistic” accuracy that is differentiable.<br>
Nonetheless, a lot of experience shows that cross entropy has better “trainability”<br>
properties than such a probabilistic accuracy, so cross entropy is generally the<br>
go-to loss for training classifiers.</p>
<aside class="quote no-group">
<blockquote>
<p>I have the feeling that my DiceLoss isn’t really working as intended.</p>
</blockquote>
</aside>
<p>It does appear that your Dice-loss code has some issues – see below.</p>
<aside class="quote no-group">
<blockquote>
<p>it would be interesting for me to utilize different loss functions and see how they compare.</p>
</blockquote>
</aside>
<p>That’s always worth doing. I would recommend comparing <code>CrossEntropyLoss</code><br>
(with class weights), some version of your Dice loss, and a combination of the two.</p>
<aside class="quote no-group">
<blockquote>
<p>Okay, so a combination of Diceloss and CrossEntropyLoss? Would you just sum up the loss?</p>
</blockquote>
</aside>
<p>Yes (but with an adjustable weight).  Something like:</p>
<pre data-code-wrap="python"><code class="lang-python">total_loss = cross_entropy_loss + dice_weight * dice_loss
total_loss.backward()
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>And if so, would you still apply class weights to CrossEntropyLoss?</p>
</blockquote>
</aside>
<p>Yes, absolutely (assuming that there is a significant imbalance in the pixel classes,<br>
which there likely is).</p>
<p>Some comments about the Dice-loss code from your first post:</p>
<aside class="quote no-group" data-username="Tony_W" data-post="1" data-topic="202228">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tony_w/48/65459_2.png" class="avatar"> Tony_W:</div>
<blockquote>
<pre><code class="lang-auto">            intersection_class_specific = (targets_one_hot_class_specific * probabilities_class_specific).sum()

            mod_a = intersection_class_specific.sum()
</code></pre>
</blockquote>
</aside>
<p>This looks unnecessary or like a typo.  <code>intersection_class_specific</code> has already<br>
been summed over pixels, so the <code>.sum()</code> in the expression for <code>mod_a</code> does nothing.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">            mod_b = targets_one_hot_class_specific.numel()
</code></pre>
</blockquote>
</aside>
<p>This is wrong.  As written, <code>mod_b</code> is just the number of pixels and is a constant,<br>
independent of <code>n_class</code>.  You want something like:</p>
<pre data-code-wrap="python"><code class="lang-python">mod_b = targets_one_hot_class_specific.sum()
</code></pre>
<p>which counts the number of pixels that are in class <code>n_class</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        dice_loss_average=-dice_coeff_average.log()
</code></pre>
</blockquote>
</aside>
<p>I do not think that taking the <code>log()</code> of the Dice coefficient is unreasonable.  However,<br>
as originally proposed, Dice loss involved maximizing (a probabilistic version of) the<br>
Dice coefficient (which is often structured as minimizing <code>1 - Dice_coefficient</code> or<br>
just <code>-Dice_coefficient</code>), so you might try that.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Tony_W
cooked: <p>Hi K. Frank, thanks again for taking the time to respond to my question and for the insights.</p>
<blockquote>
<pre><code class="lang-auto">total_loss = cross_entropy_loss + dice_weight * dice_loss
total_loss.backward()
</code></pre>
</blockquote>
<p>I will definitely try this out and report whether this improved the accuracy or not.</p>
<blockquote>
<p>This looks unnecessary or like a typo. <code>intersection_class_specific</code> has already<br>
been summed over pixels, so the <code>.sum()</code> in the expression for <code>mod_a</code> does nothing.</p>
</blockquote>
<p>Thats for pointing that out!</p>
<blockquote>
<pre><code class="lang-auto">  mod_b = targets_one_hot_class_specific.numel()
</code></pre>
</blockquote>
<p>Regarding this, I copied the basic formula from another post that I saw. I realized that this wasn’t the actual formula for DiceLoss but thought that since the <code>-log</code> was applied to the function, it would be beneficial to take a relatively small number, since the <code>-log</code> makes it bigger.</p>
<blockquote>
<p>I do not think that taking the <code>log()</code> of the Dice coefficient is unreasonable. However,<br>
as originally proposed, Dice loss involved maximizing (a probabilistic version of) the<br>
Dice coefficient (which is often structured as minimizing <code>1 - Dice_coefficient</code> or<br>
just <code>-Dice_coefficient</code>), so you might try that.</p>
</blockquote>
<p>Alright, thanks I’ll try to out too. Here, I also thought that since CrossEntropy also utilizes a log function, it would be easier during training to minimize that value (compared to a linear function).</p>
<p>Best regards, Tony!</p>

username: KFrank
cooked: <p>Hi Tony!</p>
<aside class="quote no-group" data-username="Tony_W" data-post="5" data-topic="202228" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tony_w/48/65459_2.png" class="avatar"> Tony_W:</div>
<blockquote>
<p>I also thought that since CrossEntropy also utilizes a log function, it would be easier during training to minimize that value (compared to a linear function).</p>
</blockquote>
</aside>
<p>I think this might be true (but I don’t know that it is).</p>
<p>I don’t have a sound theoretical explanation for why <code>CrossEntropyLoss</code> works<br>
as well as it does.</p>
<p>However, when the predicted probability for the correct class goes to zero,<br>
<code>CrossEntropyLoss</code> has a logarithmic divergence.  My intuition has always been<br>
that this logarithmic divergence is very helpful for training.  (On the other hand,<br>
I’ve added logarithmic divergences to other loss functions “by hand,” and doing<br>
so hasn’t helped as much as I might have hoped.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: myuser
cooked: <p>Hi Tony,</p>
<p>did the <code>CrossEntropyLoss</code> worked better in the end?</p>