Implementation of dice loss

Nishanth_Sasankan: <p>Hi All,<br>
I am trying to implement dice loss for semantic segmentation using FCN_resnet101. For some reason, the dice loss is not changing and the model is not updated.</p>
<pre><code class="lang-auto">import torch
import torchvision
import loader
from loader import DataLoaderSegmentation
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data.sampler import SubsetRandomSampler
from torch.autograd import Variable

batch_size = 1
validation_split = .2
shuffle_dataset = True
random_seed= 66

n_class    = 2
num_epochs = 5
lr         = 1e-4
momentum   = 0.9
w_decay    = 1e-5
step_size  = 50
gamma      = 0.5
cw = [0.5028138887905369,89.34501791284346]
class_weight=torch.FloatTensor(cw).cuda()
traindata = DataLoaderSegmentation('/home/ubuntu/Downloads/imgs/lensonly')
trainloader = torch.utils.data.DataLoader(traindata)

class diceloss(torch.nn.Module):
    def init(self):
        super(diceLoss, self).init()
    def forward(self,pred, target):
       smooth = 1.
       iflat = pred.contiguous().view(-1)
       tflat = target.contiguous().view(-1)
       intersection = (iflat * tflat).sum()
       A_sum = torch.sum(iflat * iflat)
       B_sum = torch.sum(tflat * tflat)
       return 1 - ((2. * intersection + smooth) / (A_sum + B_sum + smooth) )

model =torchvision.models.segmentation.fcn_resnet101(pretrained=False,aux_loss=None,num_classes=2).cuda()
criterion = diceloss()

for iter in range(num_epochs):
    for (i,l) in trainloader:
        optimizer.zero_grad()
        i= i.to(device)
        l = l.to(device)
        l[l!=0]=1
        l=l.long()
        outt = model(i)
        pred = Variable(torch.argmax(outt['out'],1).float(),requires_grad=True).cuda()
        loss = criterion(pred,l.squeeze(0).type(torch.cuda.FloatTensor))
        loss.backward()
        optimizer.step()
    print(iter)
torch.save(model, '/home/ubuntu/Downloads/newnet.pth')
</code></pre>
<p>Any and all help is appreciated. I copy pasted dice loss from here<br>
<aside class="onebox githubgist">
  <header class="source">
      <a href="https://gist.github.com/weiliu620/52d140b22685cf9552da4899e2160183" target="_blank" rel="nofollow noopener">gist.github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://gist.github.com/weiliu620/52d140b22685cf9552da4899e2160183" target="_blank" rel="nofollow noopener">https://gist.github.com/weiliu620/52d140b22685cf9552da4899e2160183</a></h4>
<h5>Dice_coeff_loss.py</h5>
<pre><code class="Python">def dice_loss(pred, target):
    """This definition generalize to real valued pred and target vector.
This should be differentiable.

    pred: tensor with first dimension as batch
    target: tensor with first dimension as batch

    """

    smooth = 1.</code></pre>
This file has been truncated. <a href="https://gist.github.com/weiliu620/52d140b22685cf9552da4899e2160183" target="_blank" rel="nofollow noopener">show original</a>

<p>
</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
Please let me know if this is the correct way of implementing it.</p>
<p>Thanks<br>
Nishanth</p>

SoucheChapich: <p>Hello,<br>
Your loss seems correct, have you different classes on your dataset?</p>
<p>Using argmax is probably not a good idea since it is not differentiable, you maybe should try to encode your labels in a one hot vector to overcome that problem.</p>

SoucheChapich: <p>The output of your network should be something like NxCxHxW and your target NxHxW where N is the batch size and C is the number of classes.</p>

Nishanth_Sasankan: <p>Thanks for your reply, I am performing a binary segmentation where I have 2 classes. I am pretty new to this platform and deep learning in general. Is there a reference which I can look up to see how to perform the following</p>
<p>“you maybe should try to encode your labels in a one hot vector to overcome that problem.”</p>
<p>any help is greatly appreciated.<br>
Thanks<br>
Nishanth</p>

SoucheChapich: <p>I am not sure but for me the problem is the fact that argmax is not a differentiable operation, so you have to overcome that.</p>
<p>If I resume, the output of your network is Nx2xHxW right ? Your idea is to take the argument max of the 2 classes and create your prediction with that information because your target is only NxHxW. The idea is to transform your target into Nx2xHxW in order to match the output dimension and compute the dice loss without applying any argmax. To transform your target from NxHxW into Nx2xHxW you can transform it to a one-hot vector like:</p>
<pre><code class="lang-auto">labels = F.one_hot(labels, num_classes = nb_classes).permute(0,3,1,2).contiguous() #in your case nb_classes = 2
</code></pre>
<p>Basically, one_hot will put 1 where the class is and 0 otherwise. (you can check the doc it is more clear than my explanation) The output of one_hot is NxHxWxC so you have to permute to obtain NxCxHxW.</p>
<p>After this transformation on your target, you should be able to compute the Dice Loss. (so you have also to make small changes in the forward of your loss)</p>

Nishanth_Sasankan: <p>Thanks for your help, so will I be calculating dice loss for each channel “C” in the NxCxHxW? if so how is the final dice loss calculated?</p>

SoucheChapich: <p>Yes exactly, you will compute the “dice loss” for every channel “C”.</p>
<p>The final loss could then be calculated as the weighted sum of all the “dice loss”.<br>
Something like :</p>
<blockquote>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/4/49b73df4420a2c43912190e1038f3ba62c01007f.png" alt="image" data-base62-sha1="aw7rS0MgS89obDKKNgwM4WIAc3R" width="68" height="51"></p>
</blockquote>
<p>where c = 2 for your case and wi is the weight you want to give at class i and Dc is like your diceloss that you linked but slightly modificated to handle one hot etc <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

dhruvbird: <p>Strangely, I haven’t found a differentiable multi-class dice loss online. Even the one on <a href="https://torchmetrics.readthedocs.io/en/stable/classification/dice.html" rel="noopener nofollow ugc">torchmetrics</a> seems to not be a differentiable Dice coefficient metric.</p>
<p>I came up with this code for using Dice as a multi-class loss. Posting here in case anyone else is stuck, and may find it useful.</p>
<pre><code class="lang-auto">class MulticlassDiceLoss(nn.Module):
    """Reference: https://www.kaggle.com/code/bigironsphere/loss-function-library-keras-pytorch#Dice-Loss
    """
    def __init__(self, num_classes, softmax_dim=None):
        super().__init__()
        self.num_classes = num_classes
        self.softmax_dim = softmax_dim
    def forward(self, logits, targets, reduction='mean', smooth=1e-6):
        """The "reduction" argument is ignored. This method computes the dice
        loss for all classes and provides an overall weighted loss.
        """
        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)
        print(targets_one_hot.shape)
        # Convert from NHWC to NCHW
        targets_one_hot = targets_one_hot.permute(0, 3, 1, 2)
        
        # Multiply one-hot encoded ground truth labels with the probabilities to get the
        # prredicted probability for the actual class.
        intersection = (targets_one_hot * probabilities).sum()
        
        mod_a = intersection.sum()
        mod_b = targets.numel()
        
        dice_coefficient = 2. * intersection / (mod_a + mod_b + smooth)
        dice_loss = -dice_coefficient.log()
        return dice_loss
# end class MulticlassDiceLoss

criterion = MulticlassDiceLoss(num_classes=3, softmax_dim=1)
y = torch.randn(10, 3, 4, 4)
ground_truth = torch.randint(0, 3, (10, 4, 4))
print(y.shape)
criterion(y, ground_truth)
</code></pre>

osman_bayram: <p>How do you test for differentiability of your loss class?</p>

dhruvbird: <p>The loss function is differentiable if the individual components are piecewise differentiable. Typically non-differentiable operations that cause the loss function to be non-differentiable include (but are not limited to).</p>
<ol>
<li>max</li>
<li>argmax</li>
<li>clip</li>
</ol>
<p>The function above uses ADD, MUL, LOG, DIV, which are all differentiable operations.</p>

Tony_W: <p>Hey <a class="mention" href="/u/dhruvbird">@dhruvbird</a> , why are you using targets.numel() and not targets.sum()?</p>
<p>Is it because the loss is supposed to be minimized (combination of .numel() and negative log())?</p>