Loss functions

Segmentation Losses

DiceLoss

class DiceLoss(include_background=True, do_sigmoid=False, do_softmax=False)[source]

Multiclass dice loss. Input logits pred (BNHW[D] where N is number of classes) is compared with ground truth ground’ (B1HW[D]). Axis N of `pred is expected to have logit predictions for each class rather than being image channels, while the same axis of ground should be 1. If the N channel of pred is 1 binary dice loss will be calculated. The smooth parameter is a value added to the intersection and union components of the inter-over-union calculation to smooth results and prevent divide-by-0, this value should be small. The include_background class attribute can be set to False for an instance of DiceLoss to exclude the first category (channel index 0) which is by convention assumed to be background. If the non-background segmentations are small compared to the total image size they can get overwhelmed by the signal from the background so excluding it in such cases helps convergence.

Parameters
  • include_background (bool) – If False channel index 0 (background category) is excluded from the calculation.

  • do_sigmoid (bool) – If True, apply a sigmoid function to the prediction.

  • do_softmax (bool) – If True, apply a softmax function to the prediction.

forward(pred, ground, smooth=1e-05)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

GeneralizedDiceLoss

class GeneralizedDiceLoss(include_background=True, do_sigmoid=False, do_softmax=False, w_type='square')[source]

Compute the generalised Dice loss defined in:

Sudre, C. et. al. (2017) Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations. DLMIA 2017.

Adapted from:

https://github.com/NifTK/NiftyNet/blob/v0.6.0/niftynet/layer/loss_segmentation.py#L279

Parameters
  • include_background (bool) – If False channel index 0 (background category) is excluded from the calculation.

  • do_sigmoid (bool) – If True, apply a sigmoid function to the prediction.

  • do_softmax (bool) – If True, apply a softmax function to the prediction.

  • w_type ('square'|'simple'|'uniform') – type of function to transform ground truth volume to a weight factor.

forward(pred, ground, smooth=1e-05)[source]
Parameters
  • pred (tensor) – the shape should be BNH[WD].

  • ground (tensor) – the shape should be B1H[WD].

  • smooth (float) – a small constant to avoid nan.