Repository: mlcommons/GaNDLF

Issue:
Improve documentation for loss functions and reduce redundancy
**Is your feature request related to a problem? Please describe.**
Some of the loss functions in the GaNDLF lack appropriate documentation and have inconsistent input variable names that do not follow the guidelines outlined in the [readme.md](https://github.com/mlcommons/GaNDLF/blob/master/GANDLF/losses/README.md). Additionally, during runtime, certain loss functions such as MCD may execute some inefficient calculations, resulting in suboptimal performance. 

**Describe the solution you'd like**
1. Inconsistent input variable names: Review the input variable names across all loss functions and ensure they are consistent.
Update the [readme.md](https://github.com/mlcommons/GaNDLF/blob/master/GANDLF/losses/README.md) to reflect the new input variable names and guidelines for future contributions.
2. Lack of documentation: Provide detailed documentation for each loss function, including a description of the function, its parameters, and its output.
3. Inefficient calculations: Review the code for each loss function and identify any inefficient calculations that can be optimized.
Optimize the code to reduce the runtime of the loss function while still producing the same results. Test the optimized code to ensure that it still produces accurate results.

**Describe alternatives you've considered**

**Additional context**



Discussion:

Resolving patch:
diff --git a/GANDLF/losses/hybrid.py b/GANDLF/losses/hybrid.py
--- a/GANDLF/losses/hybrid.py
+++ b/GANDLF/losses/hybrid.py
@@ -2,11 +2,49 @@
 from .regression import CCE_Generic, CE, CE_Logits
 
 
-def DCCE(pm, gt, params):
-    dcce_loss = MCD_loss(pm, gt, params) + CCE_Generic(pm, gt, params, CE)
+def DCCE(predicted_mask, ground_truth, params):
+    """
+    Calculates the Dice-Cross-Entropy loss.
+
+    Parameters
+    ----------
+    predicted_mask : torch.Tensor
+        Predicted mask
+    ground_truth : torch.Tensor
+        Ground truth mask
+    params : dict
+        Dictionary of parameters
+
+    Returns
+    -------
+    torch.Tensor
+        Calculated loss
+    """
+    dcce_loss = MCD_loss(predicted_mask, ground_truth, params) + CCE_Generic(
+        predicted_mask, ground_truth, params, CE
+    )
     return dcce_loss
 
 
-def DCCE_Logits(pm, gt, params):
-    dcce_loss = MCD_loss(pm, gt, params) + CCE_Generic(pm, gt, params, CE_Logits)
+def DCCE_Logits(predicted_mask, ground_truth, params):
+    """
+    Calculates the Dice-Cross-Entropy loss using logits.
+
+    Parameters
+    ----------
+    predicted_mask : torch.Tensor
+        Predicted mask logits
+    ground_truth : torch.Tensor
+        Ground truth mask
+    params : dict
+        Dictionary of parameters
+
+    Returns
+    -------
+    torch.Tensor
+        Calculated loss
+    """
+    dcce_loss = MCD_loss(predicted_mask, ground_truth, params) + CCE_Generic(
+        predicted_mask, ground_truth, params, CE_Logits
+    )
     return dcce_loss
diff --git a/GANDLF/losses/regression.py b/GANDLF/losses/regression.py
--- a/GANDLF/losses/regression.py
+++ b/GANDLF/losses/regression.py
@@ -1,44 +1,84 @@
 import torch
+import torch.nn.functional as F
 from torch.nn import MSELoss, CrossEntropyLoss, L1Loss
 from GANDLF.utils import one_hot
 
 
-def CEL(out, target, params):
+def CEL(prediction, target, params):
+    """
+    Cross entropy loss with optional class weights.
+
+    Args:
+        prediction (torch.Tensor): prediction tensor from the model.
+        target (torch.Tensor): Target tensor of class targets.
+        params (dict): Dictionary of parameters including weights.
+
+    Returns:
+        torch.Tensor: Cross entropy loss tensor.
+    """
     if len(target.shape) > 1 and target.shape[-1] == 1:
         target = torch.squeeze(target, -1)
 
     weights = None
-    if params["weights"]:
+    if params.get("weights") is not None:
+        # Check that the number of classes matches the number of weights
+        num_classes = len(params["weights"])
+        assert prediction.shape[-1] == num_classes, f"Number of classes {num_classes} does not match prediction shape {prediction.shape[-1]}"
+
         weights = torch.FloatTensor(list(params["weights"].values()))
         weights = weights.float().to(target.device)
 
     cel = CrossEntropyLoss(weight=weights)
-    return cel(out, target)
+    return cel(prediction, target)
 
 
-def CE_Logits(out, target):
-    iflat = out.contiguous().view(-1)
-    tflat = target.contiguous().view(-1)
+def CE_Logits(prediction, target):
+    """
+    Binary cross entropy loss with logits.
+
+    Args:
+        prediction (torch.Tensor): Prediction tensor from the model.
+        target (torch.Tensor): Target tensor of binary targets.
+
+    Returns:
+        torch.Tensor: Binary cross entropy loss tensor.
+    """
+    assert torch.all(target.byte() == target), "Target tensor must be binary (0 or 1)"
+
     loss = torch.nn.BCEWithLogitsLoss()
-    loss_val = loss(iflat, tflat)
+    loss_val = loss(prediction.contiguous().view(-1), target.contiguous().view(-1))
+
     return loss_val
 
 
-def CE(out, target):
-    iflat = out.contiguous().view(-1)
-    tflat = target.contiguous().view(-1)
+def CE(prediction, target):
+    """
+    Binary cross entropy loss.
+
+    Args:
+        prediction (torch.Tensor): Prediction tensor from the model.
+        target (torch.Tensor): Target tensor of binary targets.
+
+    Returns:
+        torch.Tensor: Binary cross entropy loss tensor.
+    """
+    assert torch.all(target.byte() == target), "Target tensor must be binary (0 or 1)"
+
     loss = torch.nn.BCELoss()
-    loss_val = loss(iflat.float(), tflat.float())
+    loss_val = loss(
+        prediction.contiguous().view(-1).float(), target.contiguous().view(-1).float()
+    )
+
     return loss_val
 
 
-def CCE_Generic(out, target, params, CCE_Type):
+def CCE_Generic(prediction, target, params, CCE_Type):
     """
     Generic function to calculate CCE loss
 
     Args:
-        out (torch.tensor): The predicted output value for each pixel. dimension: [batch, class, x, y, z].
-        target (torch.tensor): The ground truth label for each pixel. dimension: [batch, class, x, y, z] factorial_class_list.
+        prediction (torch.tensor): The predicted output value for each pixel. dimension: [batch, class, x, y, z].
+        target (torch.tensor): The ground truth target for each pixel. dimension: [batch, class, x, y, z] factorial_class_list.
         params (dict): The parameter dictionary.
         CCE_Type (torch.nn): The CE loss function type.
 
@@ -47,154 +87,162 @@ def CCE_Generic(out, target, params, CCE_Type):
     """
 
     acc_ce_loss = 0
-    target = one_hot(target, params["model"]["class_list"]).type(out.dtype)
+    target = one_hot(target, params["model"]["class_list"]).type(prediction.dtype)
+
     for i in range(0, len(params["model"]["class_list"])):
-        curr_ce_loss = CCE_Type(out[:, i, ...], target[:, i, ...])
+        curr_ce_loss = CCE_Type(prediction[:, i, ...], target[:, i, ...])
         if params["weights"] is not None:
             curr_ce_loss = curr_ce_loss * params["weights"][i]
         acc_ce_loss += curr_ce_loss
+
+    # Take the mean of the loss if weights are not provided.
     if params["weights"] is None:
-        acc_ce_loss /= len(params["model"]["class_list"])
+        acc_ce_loss = torch.mean(acc_ce_loss)
+
     return acc_ce_loss
 
 
-def L1(output, label, reduction="mean", scaling_factor=1):
+def L1(prediction, target, reduction="mean", scaling_factor=1):
     """
-    Calculate the mean square error between the output variable from the network and the target
-
+    Calculate the mean absolute error between the output variable from the network and the target
     Parameters
     ----------
-    output : torch.Tensor
-        The output generated usually by the network
+    prediction : torch.Tensor
+        The prediction generated by the network
     target : torch.Tensor
-        The label for the corresponding Tensor for which the output was generated
-    reduction : string, optional
-        DESCRIPTION. The default is 'mean'.
-    scaling_factor : integer, optional
-        The scaling factor to multiply the label with
-
+        The target for the corresponding Tensor for which the output was generated
+    reduction : str, optional
+        The type of reduction to apply to the output. Can be "none", "mean", or "sum". Default is "mean".
+    scaling_factor : int, optional
+        The scaling factor to multiply the target with. Default is 1.
     Returns
     -------
     loss : torch.Tensor
-        Computed Mean Squared Error loss for the output and label
-
+        The computed Mean Absolute Error (L1) loss for the output and target
     """
-    scaling_factor = torch.as_tensor(scaling_factor)
-    label = label.float()
-    label = label * scaling_factor
-    loss_fn = L1Loss(reduction=reduction)
-    iflat = output.contiguous().view(-1)
-    tflat = label.contiguous().view(-1)
-    loss = loss_fn(iflat, tflat)
+    scaling_factor = torch.as_tensor(
+        scaling_factor, dtype=target.dtype, device=target.device
+    )
+    target = target.float() * scaling_factor
+    loss = F.l1_loss(prediction, target, reduction=reduction)
     return loss
 
 
-def L1_loss(inp, target, params):
+def L1_loss(prediction, target, params):
+    """
+    Computes the L1 loss between the predictionut tensor and the target tensor.
+
+    Parameters:
+        prediction (torch.Tensor): The predictionut tensor.
+        target (torch.Tensor): The target tensor.
+        params (dict, optional): A dictionary of hyperparameters. Defaults to None.
+
+    Returns:
+        loss (torch.Tensor): The computed L1 loss.
+    """
     acc_mse_loss = 0
-    # if inp.shape != target.shape:
-    #     sys.exit('Input and target shapes are inconsistent')
 
-    if inp.shape[0] == 1:
+    if prediction.shape[0] == 1:
         if params is not None:
             acc_mse_loss += L1(
-                inp,
+                prediction,
                 target,
                 reduction=params["loss_function"]["l1"]["reduction"],
                 scaling_factor=params["scaling_factor"],
             )
         else:
-            acc_mse_loss += L1(inp, target)
-        # for i in range(0, params["model"]["num_classes"]):
-        #    acc_mse_loss += MSE(inp[i], target[i], reduction=params["loss_function"]['mse']["reduction"])
+            acc_mse_loss += L1(prediction, target)
+
+    # Compute the L1 loss
     else:
         if params is not None:
             for i in range(0, params["model"]["num_classes"]):
                 acc_mse_loss += L1(
-                    inp[:, i, ...],
+                    prediction[:, i, ...],
                     target[:, i, ...],
                     reduction=params["loss_function"]["mse"]["reduction"],
                     scaling_factor=params["scaling_factor"],
                 )
         else:
-            for i in range(0, inp.shape[1]):
-                acc_mse_loss += L1(inp[:, i, ...], target[:, i, ...])
+            for i in range(0, prediction.shape[1]):
+                acc_mse_loss += L1(prediction[:, i, ...], target[:, i, ...])
+
+    # Normalize the loss by the number of classes
     if params is not None:
         acc_mse_loss /= params["model"]["num_classes"]
     else:
-        acc_mse_loss /= inp.shape[1]
+        acc_mse_loss /= prediction.shape[1]
 
     return acc_mse_loss
 
 
-def MSE(output, label, reduction="mean", scaling_factor=1):
+def MSE(prediction, target, reduction="mean", scaling_factor=1):
     """
     Calculate the mean square error between the output variable from the network and the target
-
     Parameters
     ----------
-    output : torch.Tensor
-        The output generated usually by the network
+    prediction : torch.Tensor
+        The prediction generated usually by the network
     target : torch.Tensor
-        The label for the corresponding Tensor for which the output was generated
+        The target for the corresponding Tensor for which the output was generated
     reduction : string, optional
         DESCRIPTION. The default is 'mean'.
-    scaling_factor : integer, optional
-        The scaling factor to multiply the label with
-
+    scaling_factor : float, optional
+        The scaling factor to multiply the target with
     Returns
     -------
     loss : torch.Tensor
-        Computed Mean Squared Error loss for the output and label
-
+        Computed Mean Squared Error loss for the output and target
     """
-    scaling_factor = torch.as_tensor(scaling_factor)
-    label = label.float()
-    label = label * scaling_factor
-    loss_fn = MSELoss(reduction=reduction)
-    iflat = output.contiguous().view(-1)
-    tflat = label.contiguous().view(-1)
-    loss = loss_fn(iflat, tflat)
+    scaling_factor = torch.as_tensor(scaling_factor, dtype=torch.float32)
+    target = target.float() * scaling_factor
+    loss = F.mse_loss(prediction, target, reduction=reduction)
     return loss
 
 
-def MSE_loss(inp, target, params):
-    acc_mse_loss = 0
-    # if inp.shape != target.shape:
-    #     sys.exit('Input and target shapes are inconsistent')
+def MSE_loss(prediction, target, params=None):
+    """
+    Compute the mean squared error loss for the predictionut and target
 
-    reduction = "mean"
-    if params is not None:
-        if "mse" in params["loss_function"]:
-            if isinstance(params["loss_function"]["mse"], dict):
-                reduction = params["loss_function"]["mse"]["reduction"]
+    Parameters
+    ----------
+    prediction : torch.Tensor
+        The predictionut tensor
+    target : torch.Tensor
+        The target tensor
+    params : dict, optional
+        A dictionary of parameters. Default: None.
+        If params is not None and contains the key "loss_function", the value of
+        "loss_function" is expected to be a dictionary with a key "mse", which
+        can contain the key "reduction" and/or "scaling_factor". If "reduction" is
+        not specified, the default is 'mean'. If "scaling_factor" is not specified,
+        the default is 1.
 
-    if inp.shape[0] == 1:
-        if params is not None:
+    Returns
+    -------
+    acc_mse_loss : torch.Tensor
+        Computed mean squared error loss for the predictionut and target
+    """
+    reduction = "mean"
+    scaling_factor = 1
+    if params is not None and "loss_function" in params:
+        mse_params = params["loss_function"].get("mse", {})
+        reduction = mse_params.get("reduction", "mean")
+        scaling_factor = mse_params.get("scaling_factor", 1)
+
+    if prediction.shape[0] == 1:
+        acc_mse_loss = MSE(
+            prediction, target, reduction=reduction, scaling_factor=scaling_factor
+        )
+    else:
+        acc_mse_loss = 0
+        for i in range(prediction.shape[1]):
             acc_mse_loss += MSE(
-                inp,
-                target,
+                prediction[:, i, ...],
+                target[:, i, ...],
                 reduction=reduction,
-                scaling_factor=params["scaling_factor"],
+                scaling_factor=scaling_factor,
             )
-        else:
-            acc_mse_loss += MSE(inp, target)
-        # for i in range(0, params["model"]["num_classes"]):
-        #    acc_mse_loss += MSE(inp[i], target[i], reduction=params["loss_function"]['mse']["reduction"])
-    else:
-        if params is not None:
-            for i in range(0, params["model"]["num_classes"]):
-                acc_mse_loss += MSE(
-                    inp[:, i, ...],
-                    target[:, i, ...],
-                    reduction=reduction,
-                    scaling_factor=params["scaling_factor"],
-                )
-        else:
-            for i in range(0, inp.shape[1]):
-                acc_mse_loss += MSE(inp[:, i, ...], target[:, i, ...])
-    if params is not None:
-        acc_mse_loss /= params["model"]["num_classes"]
-    else:
-        acc_mse_loss /= inp.shape[1]
+        acc_mse_loss /= prediction.shape[1]
 
     return acc_mse_loss
diff --git a/GANDLF/losses/segmentation.py b/GANDLF/losses/segmentation.py
--- a/GANDLF/losses/segmentation.py
+++ b/GANDLF/losses/segmentation.py
@@ -2,16 +2,16 @@
 
 
 # Dice scores and dice losses
-def dice(output, label):
+def dice(predicted, target):
     """
     This function computes a dice score between two tensors
 
     Parameters
     ----------
-    output : Tensor
-        Output predicted generally by the network
-    label : Tensor
-        Required target label to match the output with
+    predicted : Tensor
+        predicted value by the network
+    target : Tensor
+        Required target label to match the predicted with
 
     Returns
     -------
@@ -20,78 +20,193 @@ def dice(output, label):
 
     """
     smooth = 1e-7
-    iflat = output.contiguous().view(-1)
-    tflat = label.contiguous().view(-1)
-    intersection = (iflat * tflat).sum()
-    return (2.0 * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth)
 
+    predicted_flat = predicted.contiguous().view(-1)
+    label_flat = target.contiguous().view(-1)
+    intersection = (predicted_flat * label_flat).sum()
 
-def MCD(pm, gt, num_class, weights=None, ignore_class=None, loss_type=0):
+    dice_score = (2.0 * intersection + smooth) / (
+        predicted_flat.sum() + label_flat.sum() + smooth
+    )
+
+    return dice_score
+
+
+def MCD(predicted, target, num_class, weights=None, ignore_class=None, loss_type=0):
     """
-    These weights should be the dice weights, not dice weights
-    loss_type:
+    Parameters
+    ----------
+    predicted : torch.Tensor
+        Prediction tensor generated by the model
+    tarrget : torch.Tensor
+        Ground truth tensor
+    num_class : int
+        Number of classes (including the background class)
+    weights : list, optional
+        Dice weights for each class (excluding the background class), defaults to None
+    ignore_class : int, optional
+        Class to ignore, defaults to None
+    loss_type : int, optional
+        Type of loss to compute, defaults to 0
         0: no loss, normal dice calculation
         1: dice loss, (1-dice)
         2: log dice, -log(dice)
+
+    Returns
+    -------
+    Tensor
+        Mean Class Dice score
     """
+
     acc_dice = 0
-    for i in range(0, num_class):  # 0 is background
-        currentDice = dice(gt[:, i, ...], pm[:, i, ...])
+
+    for i in range(num_class):  # 0 is background
+
+        currentDice = dice(predicted[:, i, ...], target[:, i, ...])
+
         if loss_type == 1:
             currentDice = 1 - currentDice  # subtract from 1 because this is a loss
         elif loss_type == 2:
             # negative because we want positive losses
-            currentDice = -torch.log(currentDice + torch.finfo(torch.float32).eps)
+            currentDice = -torch.log(
+                currentDice + torch.finfo(torch.float32).eps
+            )
+
         if weights is not None:
-            currentDice = currentDice * weights[i]
+            currentDice = currentDice * weights[i]  # multiply by weight
+
         acc_dice += currentDice
+
     if weights is None:
         acc_dice /= num_class  # we should not be considering 0
+
     return acc_dice
 
 
-def MCD_loss(pm, gt, params):
+def MCD_loss(predicted, target, params):
     """
     These weights should be the penalty weights, not dice weights
     """
-    return MCD(pm, gt, len(params["model"]["class_list"]), params["weights"], None, 1)
+    return MCD(
+        predicted,
+        target,
+        len(params["model"]["class_list"]),
+        params["weights"],
+        None,
+        1,
+    )
 
 
-def MCD_log_loss(pm, gt, params):
+def MCD_log_loss(predicted, target, params):
     """
     These weights should be the penalty weights, not dice weights
     """
-    return MCD(pm, gt, len(params["model"]["class_list"]), params["weights"], None, 2)
+    return MCD(
+        predicted,
+        target,
+        len(params["model"]["class_list"]),
+        params["weights"],
+        None,
+        2,
+    )
 
 
-def tversky(inp, target, alpha):
-    smooth = 1e-7
-    iflat = inp.view(-1)
-    tflat = target.view(-1)
-    intersection = (iflat * tflat).sum()
-    fps = (iflat * (1 - tflat)).sum()
-    fns = ((1 - iflat) * tflat).sum()
-    denominator = intersection + (alpha * fps) + ((1 - alpha) * fns) + smooth
-    return (intersection + smooth) / denominator
+def tversky_loss(predicted, target, alpha=0.5, beta=0.5, smooth=1e-7):
+    """
+    This function calculates the Tversky loss between two tensors.
 
+    Parameters
+    ----------
+    predicted : torch.Tensor
+        predicted predicted generally by the network
+    target : torch.Tensor
+        Required target label to match the predicted with
+    alpha : float, optional
+        Weight of false positives
+    beta : float, optional
+        Weight of false negatives
+    smooth : float, optional
+        Smoothing factor to avoid division by zero
+
+    Returns
+    -------
+    torch.Tensor
+        Computed Tversky Loss
+
+    """
+    # Move this part later to parameter parsing, no need to check every time
+    assert 0 <= alpha <= 1, f"Invalid alpha value: {alpha}"
+    assert 0 <= beta <= 1, f"Invalid beta value: {beta}"
+    assert 0 <= alpha + beta <= 1, f"Invalid alpha and beta values: {alpha}, {beta}"
 
-def tversky_loss(inp, target, alpha=1):
-    tversky_val = tversky(inp, target, alpha)
-    return 1 - tversky_val
+    predicted_flat = predicted.contiguous().view(-1)
+    target_flat = target.contiguous().view(-1)
 
+    true_positives = (predicted_flat * target_flat).sum()
+    false_positives = ((1 - target_flat) * predicted_flat).sum()
+    false_negatives = (target_flat * (1 - predicted_flat)).sum()
+
+    numerator = true_positives
+    denominator = true_positives + alpha * false_positives + beta * false_negatives
+    score = (numerator + smooth) / (denominator + smooth)
+
+    loss = 1 - score
+    return loss
+
+
+def MCT_loss(predicted, target, params=None):
+    """
+    This function calculates the Multi-Class Tversky loss between two tensors.
+
+    Parameters
+    ----------
+    predicted : torch.Tensor
+        predicted predicted generally by the network
+    target : torch.Tensor
+        Required target label to match the predicted with
+    params : dict, optional
+        Additional parameters for computing loss function, including weights for each class
+
+    Returns
+    -------
+    torch.Tensor
+        Computed Multi-Class Tversky Loss
+
+    """
 
-def MCT_loss(inp, target, params):
     acc_tv_loss = 0
-    for i in range(0, len(params["model"]["class_list"])):
-        curr_loss = tversky_loss(inp[:, i, ...], target[:, i, ...])
-        if params["weights"] is not None:
+    num_classes = predicted.shape[1]
+
+    for i in range(num_classes):
+        curr_loss = tversky_loss(predicted[:, i, ...], target[:, i, ...])
+        if params is not None and params.get("weights") is not None:
             curr_loss = curr_loss * params["weights"][i]
         acc_tv_loss += curr_loss
-    if params["weights"] is None:
-        acc_tv_loss /= len(params["model"]["class_list"])
+
+    if params is not None and params.get("weights") is None:
+        acc_tv_loss /= num_classes
+
     return acc_tv_loss
 
 
 def KullbackLeiblerDivergence(mu, logvar, params=None):
+    """
+    Calculates the Kullback-Leibler divergence between two Gaussian distributions.
+
+    Parameters
+    ----------
+    mu : torch.Tensor
+        The mean of the first Gaussian distribution
+    logvar : torch.Tensor
+        The logarithm of the variance of the first Gaussian distribution
+    params : dict, optional
+        A dictionary of optional parameters
+
+    Returns
+    -------
+    torch.Tensor
+        The computed Kullback-Leibler divergence
+
+    """
     loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1)
     return loss.mean()


Test patch: