text: """Modified from https://github.com/LikeLy-Journey/SegmenTron/blob/master/
  segmentron/solver/loss.py (Apache-2.0 License)"""
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from ..builder import LOSSES
  from .utils import get_class_weight, weighted_loss
  
  
  @weighted_loss
  def dice_loss(pred,
                target,
                valid_mask,
                smooth=1,
                exponent=2,
                class_weight=None,
                ignore_index=255):
      assert pred.shape[0] == target.shape[0]
  
      total_loss = 0
      num_classes = pred.shape[1]
      for i in range(num_classes):
          if i != ignore_index:
              dice_loss_value = binary_dice_loss(
                  pred[:, i],
                  target[..., i],
                  valid_mask=valid_mask,
                  smooth=smooth,
                  exponent=exponent
              )
  
              if class_weight is not None:
                  dice_loss_value *= class_weight[i]
  
              total_loss += dice_loss_value
  
      return total_loss / num_classes
  
  
  @weighted_loss
  def binary_dice_loss(pred, target, valid_mask, smooth=1, exponent=2):
      assert pred.shape[0] == target.shape[0]
  
      pred = pred.reshape(pred.shape[0], -1)
      target = target.reshape(target.shape[0], -1)
      valid_mask = valid_mask.reshape(valid_mask.shape[0], -1)
  
      valid_pred = torch.mul(pred, valid_mask)
      valid_target = torch.mul(target, valid_mask)
  
      num = torch.sum(torch.mul(valid_pred, valid_target), dim=1) * 2 + smooth
      den = torch.sum(valid_pred.pow(exponent) + valid_target.pow(exponent), dim=1) + smooth
  
      return 1 - num / den
  
  
  @LOSSES.register_module()
  class DiceLoss(nn.Module):
      """DiceLoss.
  
      This loss is proposed in `V-Net: Fully Convolutional Neural Networks for
      Volumetric Medical Image Segmentation <https://arxiv.org/abs/1606.04797>`_.
  
      Args:
          loss_type (str, optional): Binary or multi-class loss.
              Default: 'multi_class'. Options are "binary" and "multi_class".
          smooth (float): A float number to smooth loss, and avoid NaN error.
              Default: 1
          exponent (float): An float number to calculate denominator
              value: \\sum{x^exponent} + \\sum{y^exponent}. Default: 2.
          reduction (str, optional): The method used to reduce the loss. Options
              are "none", "mean" and "sum". This parameter only works when
              per_image is True. Default: 'mean'.
          class_weight (list[float] | str, optional): Weight of each class. If in
              str format, read them from a file. Defaults to None.
          loss_weight (float, optional): Weight of the loss. Default to 1.0.
      """
  
      def __init__(self,
                   smooth=1,
                   exponent=2,
                   reduction='mean',
                   class_weight=None,
                   loss_weight=1.0):
          super(DiceLoss, self).__init__()
  
          self.smooth = smooth
          self.exponent = exponent
          self.reduction = reduction
          self.class_weight = get_class_weight(class_weight)
          self.loss_weight = loss_weight
  
      @property
      def name(self):
          return 'dice'
  
      def forward(self,
                  pred,
                  target,
                  avg_factor=None,
                  reduction_override=None,
                  ignore_index=255,
                  **kwargs):
          assert reduction_override in (None, 'none', 'mean', 'sum')
          reduction = (reduction_override if reduction_override else self.reduction)
  
          if self.class_weight is not None:
              class_weight = pred.new_tensor(self.class_weight)
          else:
              class_weight = None
  
          pred = F.softmax(pred, dim=1)
          num_classes = pred.shape[1]
          one_hot_target = F.one_hot(
              torch.clamp(target.long(), 0, num_classes - 1),
              num_classes=num_classes
          )
  
          valid_mask = (target != ignore_index).long()
  
          loss = self.loss_weight * dice_loss(
              pred,
              one_hot_target,
              valid_mask=valid_mask,
              reduction=reduction,
              avg_factor=avg_factor,
              smooth=self.smooth,
              exponent=self.exponent,
              class_weight=class_weight,
              ignore_index=ignore_index,
              **kwargs
          )
  
          return loss
  