class_docstr: Cross entropy loss for the task of semantic segmentation.
  
  Args:
      opts: command-line arguments
docstr: Compute CE segmentation loss
  
  Args:
      input_sample: Input image tensor to model.
      prediction: Output of model. It can be a
          * Tensor
          * Tuple[Tensor, Tensor]
          * Mapping[segmentation_output, Tensor]
          * Mapping[segmentation_output, Tuple[Tensor, Tensor]], where segmentation_output is a required key.
      target: Target label tensor containing values in the range `[0, C)`, where :math:`C` is the number of classes
  
  Shapes:
      input_sample: This loss function does not care about this argument.
      prediction:
          * When prediction is a Tensor, then shape is [Batch size, C, Height, Width]
          * When prediction is a Tuple[Tensor, Tensor], then shape of one tensor is [Batch size, C, Height, Width]
              while the other is [Batch size, C, Height / O, Width/ O]
              where O is the output stride of feature map (typically 4).
          * When prediction is a dictionary, then the shape of prediction["segmentation_output"] should
              be the same as described in above steps (depending on type).
      target: The shape of target tensor is [Batch size, Height, Width]
  
  Returns:
      Mapping of the form (string: scalar value) is returned with total_loss as mandatory and
      (seg_loss, aux_loss) as optional keys. total_loss is weighted sum of seg_loss and aux_loss (when applicable).
docstr_tok: ["Compute","CE","segmentation","loss","Args",":","input_sample",":","Input","image","tensor","to","model",".","prediction",":","Output","of","model",".","It","can","be","a","*","Tensor","*","Tuple","[","Tensor",",","Tensor","]","*","Mapping","[","segmentation_output",",","Tensor","]","*","Mapping","[","segmentation_output",",","Tuple","[","Tensor",",","Tensor","]","]",",","where","segmentation_output","is","a","required","key",".","target",":","Target","label","tensor","containing","values","in","the","range","`","[","0",",","C",")","`",",","where",":","math",":","`","C","`","is","the","number","of","classes","Shapes",":","input_sample",":","This","loss","function","does","not","care","about","this","argument",".","prediction",":","*","When","prediction","is","a","Tensor",",","then","shape","is","[","Batch","size",",","C",",","Height",",","Width","]","*","When","prediction","is","a","Tuple","[","Tensor",",","Tensor","]",",","then","shape","of","one","tensor","is","[","Batch","size",",","C",",","Height",",","Width","]","while","the","other","is","[","Batch","size",",","C",",","Height","\/","O",",","Width\/","O","]","where","O","is","the","output","stride","of","feature","map","(","typically","4",")",".","*","When","prediction","is","a","dictionary",",","then","the","shape","of","prediction","[","``","segmentation_output","''","]","should","be","the","same","as","described","in","above","steps","(","depending","on","type",")",".","target",":","The","shape","of","target","tensor","is","[","Batch","size",",","Height",",","Width","]","Returns",":","Mapping","of","the","form","(","string",":","scalar","value",")","is","returned","with","total_loss","as","mandatory","and","(","seg_loss",",","aux_loss",")","as","optional","keys",".","total_loss","is","weighted","sum","of","seg_loss","and","aux_loss","(","when","applicable",")","."]
code: def forward(
      self,
      input_sample: Any,
      prediction: Union[
          Mapping[str, Union[Tensor, Tuple[Tensor, Tensor]]],
          Tensor,
          Tuple[Tensor, Tensor],
      ],
      target: Tensor,
      *args,
      **kwargs,
  ) -> Mapping[str, Tensor]:
      """Compute CE segmentation loss
  
      Args:
          input_sample: Input image tensor to model.
          prediction: Output of model. It can be a
              * Tensor
              * Tuple[Tensor, Tensor]
              * Mapping[segmentation_output, Tensor]
              * Mapping[segmentation_output, Tuple[Tensor, Tensor]], where segmentation_output is a required key.
          target: Target label tensor containing values in the range `[0, C)`, where :math:`C` is the number of classes
  
      Shapes:
          input_sample: This loss function does not care about this argument.
          prediction:
              * When prediction is a Tensor, then shape is [Batch size, C, Height, Width]
              * When prediction is a Tuple[Tensor, Tensor], then shape of one tensor is [Batch size, C, Height, Width]
                  while the other is [Batch size, C, Height / O, Width/ O]
                  where O is the output stride of feature map (typically 4).
              * When prediction is a dictionary, then the shape of prediction["segmentation_output"] should
                  be the same as described in above steps (depending on type).
          target: The shape of target tensor is [Batch size, Height, Width]
  
      Returns:
          Mapping of the form (string: scalar value) is returned with total_loss as mandatory and
          (seg_loss, aux_loss) as optional keys. total_loss is weighted sum of seg_loss and aux_loss (when applicable).
      """
  
      if isinstance(prediction, (Tuple, Tensor)):
          return self._forward_seg(
              prediction=prediction, target=target, *args, **kwargs
          )
      elif isinstance(prediction, Mapping):
          if "segmentation_output" not in prediction:
              logger.error(
                  f"segmentation_output is a mandatory key in prediction when"
                  f"type of prediction is Dict. Got: {prediction.keys()}"
              )
  
          seg_loss = self._forward_seg(
              prediction=prediction["segmentation_output"],
              target=target,
              *args,
              **kwargs,
          )
  
          return seg_loss
      else:
          logger.error(
              f"Prediction should be either a Tensor or Tuple[Tensor, Tensor] "
              f"or Dictionary[str, Tensor] in {self.__class__.__name__}. Got: {type(prediction)}"
          )
  
file_path: ml-cvnets/loss_fn/segmentation/cross_entropy.py
imports: import argparse
  from typing import Any, Mapping, Optional, Tuple, Union
  from torch import Tensor
  from torch.nn import functional
  from loss_fn import LOSS_REGISTRY
  from loss_fn.segmentation.base_segmentation_criteria import BaseSegmentationCriteria
  from loss_fn.utils.class_weighting import compute_class_weights
  from utils import logger
label_desc: Use image node_id 5 for calling the SegCrossEntropy obj's underlying member method code with example usage: obj.forward(input_sample, prediction, target) and returns: self, seg_loss