Extending multi-class 2D Dice Loss to 3D

cwat_tum: <p>Hi all,</p>
<p>I’m attempting to extend a multi-class 2D dice loss implementation to 3D. Overall, my segmentations learn from a combined loss: DICE loss + Cross Entropy Loss. My 2D segmentation implementation learns from this loss, but my 3D implementation doesn’t seem to be learning anything, so I guess it’s either the dice loss implementation or the cross entropy loss calculation that is causing this.</p>
<p>Here, the 2D snippet ( credit to this discussion: <a href="https://discuss.pytorch.org/t/one-hot-encoding-with-autograd-dice-loss/9781/3" class="inline-onebox">One-hot encoding with autograd (Dice loss)</a> )</p>
<p>def _dice_loss_multichannel(output, target, weights=None, ignore_index=None):<br>
“”"</p>
<pre><code>    :param output: NxCxHxW Variable
    :param target: NxHxW LongTensor
    :param weights: C FloatTensor
    :param ignore_index: int index to ignore from loss
    :param binary: bool for binarized one chaneel(C=1) input
    :return:
    """
    eps = 0.0001
    encoded_target = output.detach() * 0

    if ignore_index is not None:
        mask = target == ignore_index
        target = target.clone()
        target[mask] = 0
        encoded_target.scatter_(1, target.unsqueeze(1), 1)
        mask = mask.unsqueeze(1).expand_as(encoded_target)
        encoded_target[mask] = 0
    else:
        encoded_target.scatter_(1, target.unsqueeze(1), 1)

    if weights is None:
        weights = 1

    intersection = output * encoded_target
    numerator = 2 * intersection.sum(0).sum(1).sum(1)
    denominator = output + encoded_target

    if ignore_index is not None:
        denominator[mask] = 0
    denominator = denominator.sum(0).sum(1).sum(1) + eps
    loss_per_channel = weights * (1 - (numerator / denominator))

    return loss_per_channel.sum() / output.size(1)
</code></pre>
<p>It seems to me that extending this to 3D should be as easy as adding a dimension and following the same process:</p>
<p>def _dice_loss_multichannel3D(output, target, weights=None, ignore_index=None):<br>
“”"<br>
Forward pass</p>
<pre><code>    :param output: NxCxDxHxW Variable
    :param target: NxDxHxW LongTensor
    :param weights: C FloatTensor
    :param ignore_index: int index to ignore from loss
    :param binary: bool for binarized one chaneel(C=1) input
    :return:
    """
    eps = 0.0001
    encoded_target = output.detach() * 0

    if ignore_index is not None:
        mask = target == ignore_index
        target = target.clone()
        target[mask] = 0
        encoded_target.scatter_(1, target.unsqueeze(1), 1)
        mask = mask.unsqueeze(1).expand_as(encoded_target)
        encoded_target[mask] = 0
    else:
        encoded_target.scatter_(1, target.unsqueeze(1), 1)

    if weights is None:
        weights = 1

    intersection = output * encoded_target
    numerator = 2 * torch.squeeze(intersection).sum(1).sum(1).sum(1)
    denominator = output + encoded_target

    if ignore_index is not None:
        denominator[mask] = 0
    denominator = torch.squeeze(denominator).sum(1).sum(1).sum(1) + eps
    loss_per_channel = weights * (1 - (numerator / denominator))

    return loss_per_channel.sum() / loss_per_channel.size(0)
</code></pre>
<p>If anyone could comment on the validity of this approach, I would appreciate it!</p>
<p>Best,</p>
<p>Chris</p>

tom: <p>This should work (and if you put triple backticks “```” at the top and bottom of your code, it will be better formatted).<br>
The general problem of 2d-&gt;3d is whether it is prohibitively expensive in terms of memory (or compute), which could be the case if the total number of voxels is very large.</p>
<p>Best regards</p>
<p>Thomas</p>

cwat_tum: <p>Hi Tom,</p>
<p>Thanks for your reply. In my implementation, I use a combined loss, dice + cross entropy. I suppose if the dice loss is not causing the problem, I wonder if the CE loss is. According to the documentation, I can call the CE loss in the below manner on volumetric outputs and ground truths. Is this call to cross entropy also okay?</p>
<pre><code class="lang-auto">class CombinedLoss3D(_Loss):
    A combination of dice  and cross entropy loss for volumetric outputs and ground truths

    def __init__(self, weight):
        super(CombinedLoss3D, self).__init__()
        self.cross_entropy_loss = nn.CrossEntropyLoss(weight=weight)
        self.dice_loss = DiceLoss() 

    def forward(self, output, ground_truth, weight=None):
        """
        Forward pass

        :param output: torch.tensor (NxCxDxHxW) Network output (logits) not normalized. 
        :param ground_truth: torch.tensor (NxDxHxW) 
        :param weight: torch.tensor (N) 
        :return: scalar
       """
        y_2 = self.dice_loss(output, ground_truth, weight)
        y_1 = self.cross_entropy_loss.forward(output, ground_truth)

        return y_1 + y_2
</code></pre>
<p>Thank you!</p>

Raja_Haseeb: <p>Hi <a class="mention" href="/u/cwat_tum">@cwat_tum</a> . I am trying to implement <code>dice loss</code> for 3D segmentation task. I tried using the loss function defined <a href="https://colab.research.google.com/github/fepegar/torchio-notebooks/blob/main/notebooks/TorchIO_tutorial.ipynb#scrollTo=oWbWAyFNTj74&amp;line=1&amp;uniqifier=1" rel="noopener nofollow ugc">here</a>.</p>
<p>However, it give me size mismatch error because in some of the in volumes, one of the class is missing.I think either his function is only for <strong>Binary</strong> case, or it requires <code>pred</code> and <code>target</code> tensors to have the same shape.</p>
<p>I was wondering if you figured it out for the <code>multi-class</code> case and can help me out. Thank you.</p>