MRI Segmentation Losses#

class atommic.collections.segmentation.losses.CrossEntropyLoss(*args: Any, **kwargs: Any)[source]#

Bases: Module

Wrapper around PyTorch’s CrossEntropyLoss to support 2D and 3D inputs.

__init__(num_samples: int = 50, ignore_index: int = -100, reduction: str = 'none', label_smoothing: float = 0.0, weight: Optional[torch.Tensor] = None)[source]#

Inits CrossEntropyLoss.

Parameters

  • num_samples (int, optional) – Number of Monte Carlo samples, by default 50

  • ignore_index (int, optional) – Index to ignore, by default -100

  • reduction (str, optional) – Reduction method, by default “none”

  • label_smoothing (float, optional) – Label smoothing, by default 0.0

  • weight (torch.Tensor, optional) – Weight for each class, by default None

forward(target: torch.Tensor, _input: torch.Tensor, pred_log_var: Optional[torch.Tensor] = None) → torch.Tensor[source]#

Forward pass of CrossEntropyLoss.

Parameters

  • target (torch.Tensor) – Target tensor. Shape: (batch_size, num_classes, *spatial_dims)

  • _input (torch.Tensor) – Prediction tensor. Shape: (batch_size, num_classes, *spatial_dims)

  • pred_log_var (torch.Tensor, optional) – Prediction log variance tensor. Shape: (batch_size, num_classes, *spatial_dims). Default is None.

Returns

Loss tensor. Shape: (batch_size, *spatial_dims)

Return type

torch.Tensor

class atommic.collections.segmentation.losses.Dice(*args: Any, **kwargs: Any)[source]#

Bases: Loss

Wrapper for monai.losses.DiceLoss to support multi-class and multi-label tasks.

Compute average Dice loss between two tensors. It can support both multi-classes and multi-labels tasks. The data input (BNHW[D] where N is number of classes) is compared with ground truth target (BNHW[D]).

Note that axis N of input is expected to be logits or probabilities for each class, if passing logits as input, must set sigmoid=True or softmax=True, or specifying other_act. And the same axis of target can be 1 or N (one-hot format).

The smooth_nr and smooth_dr parameters are values added to the intersection and union components of the inter-over-union calculation to smooth results respectively, these values should be small.

The original paper: Milletari, F. et. al. (2016) V-Net: Fully Convolutional Neural Networks forVolumetric Medical Image Segmentation, 3DV, 2016.

Examples

>>> import torch
>>> from atommic.collections.segmentation.losses.dice import Dice
>>> pred = torch.tensor([[[[0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5]]],
...                       [[[0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5]]]])
>>> target = torch.tensor([[[[0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0]]],
...                         [[[1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1]]]])
>>> dice = Dice(include_background=False, to_onehot_y=True, sigmoid=False, softmax=False)
>>> dice(pred, target)
tensor(0.5000)
__init__(include_background: bool = True, to_onehot_y: bool = False, sigmoid: bool = True, softmax: bool = False, other_act: Optional[Callable] = None, squared_pred: bool = False, jaccard: bool = False, flatten: bool = False, reduction: str = 'mean', smooth_nr: float = 1e-05, smooth_dr: float = 1e-05, batch: bool = True)[source]#

Inits Dice.

Parameters

  • include_background (bool) – whether to skip Dice computation on the first channel of the predicted output. Default is True.

  • to_onehot_y (bool) – Whether to convert y into the one-hot format. Default is False.

  • sigmoid (bool) – Whether to add sigmoid function to the input data. Default is True.

  • softmax (bool) – Whether to add softmax function to the input data. Default is False.

  • other_act (Callable) – Use this parameter if you want to apply another type of activation layer. Default is None.

  • squared_pred (bool) – Whether to square the prediction before calculating Dice. Default is False.

  • jaccard (bool) – Whether to compute Jaccard Index as a loss. Default is False.

  • flatten (bool) – Whether to flatten input data. Default is False.

  • reduction (str) – Specifies the reduction to apply to the output: ‘none’ | ‘mean’ | ‘sum’. ‘none’: no reduction will be applied. ‘mean’: the sum of the output will be divided by the number of elements in the output. ‘sum’: the output will be summed. Default is mean.

  • smooth_nr (float) – A small constant added to the numerator to avoid nan when all items are 0. Default is 1e-5.

  • smooth_dr (float) – A small constant added to the denominator to avoid nan when all items are 0. Default is 1e-5.

  • batch (bool) – If True, compute Dice loss for each batch and return a tensor with shape (batch_size,). If False, compute Dice loss for the whole batch and return a tensor with shape (1,). Default is True.

forward(target: torch.Tensor, _input: torch.Tensor) → Tuple[Union[torch.Tensor, Any], torch.Tensor][source]#

Forward pass of Dice.

Parameters
Returns

Dice loss.

Return type

torch.Tensor