# Copyright (c) MONAI Consortium # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import warnings from collections.abc import Sequence from typing import Optional import torch import torch.nn.functional as F from torch.nn.modules.loss import _Loss from monai.networks import one_hot from monai.utils import LossReduction class FocalLoss(_Loss): """ FocalLoss is an extension of BCEWithLogitsLoss that down-weights loss from high confidence correct predictions. Reimplementation of the Focal Loss described in: - ["Focal Loss for Dense Object Detection"](https://arxiv.org/abs/1708.02002), T. Lin et al., ICCV 2017 - "AnatomyNet: Deep learning for fast and fully automated whole-volume segmentation of head and neck anatomy", Zhu et al., Medical Physics 2018 Example: >>> import torch >>> from monai.losses import FocalLoss >>> from torch.nn import BCEWithLogitsLoss >>> shape = B, N, *DIMS = 2, 3, 5, 7, 11 >>> input = torch.rand(*shape) >>> target = torch.rand(*shape) >>> # Demonstrate equivalence to BCE when gamma=0 >>> fl_g0_criterion = FocalLoss(reduction='none', gamma=0) >>> fl_g0_loss = fl_g0_criterion(input, target) >>> bce_criterion = BCEWithLogitsLoss(reduction='none') >>> bce_loss = bce_criterion(input, target) >>> assert torch.allclose(fl_g0_loss, bce_loss) >>> # Demonstrate "focus" by setting gamma > 0. >>> fl_g2_criterion = FocalLoss(reduction='none', gamma=2) >>> fl_g2_loss = fl_g2_criterion(input, target) >>> # Mark easy and hard cases >>> is_easy = (target > 0.7) & (input > 0.7) >>> is_hard = (target > 0.7) & (input < 0.3) >>> easy_loss_g0 = fl_g0_loss[is_easy].mean() >>> hard_loss_g0 = fl_g0_loss[is_hard].mean() >>> easy_loss_g2 = fl_g2_loss[is_easy].mean() >>> hard_loss_g2 = fl_g2_loss[is_hard].mean() >>> # Gamma > 0 causes the loss function to "focus" on the hard >>> # cases. IE, easy cases are downweighted, so hard cases >>> # receive a higher proportion of the loss. >>> hard_to_easy_ratio_g2 = hard_loss_g2 / easy_loss_g2 >>> hard_to_easy_ratio_g0 = hard_loss_g0 / easy_loss_g0 >>> assert hard_to_easy_ratio_g2 > hard_to_easy_ratio_g0 """ def __init__( self, include_background: bool = True, to_onehot_y: bool = False, gamma: float = 2.0, alpha: float | None = None, weight: Sequence[float] | float | int | torch.Tensor | None = None, reduction: LossReduction | str = LossReduction.MEAN, use_softmax: bool = False, ) -> None: """ Args: include_background: if False, channel index 0 (background category) is excluded from the loss calculation. If False, `alpha` is invalid when using softmax. to_onehot_y: whether to convert the label `y` into the one-hot format. Defaults to False. gamma: value of the exponent gamma in the definition of the Focal loss. Defaults to 2. alpha: value of the alpha in the definition of the alpha-balanced Focal loss. The value should be in [0, 1]. Defaults to None. weight: weights to apply to the voxels of each class. If None no weights are applied. The input can be a single value (same weight for all classes), a sequence of values (the length of the sequence should be the same as the number of classes. If not ``include_background``, the number of classes should not include the background category class 0). The value/values should be no less than 0. Defaults to None. reduction: {``"none"``, ``"mean"``, ``"sum"``} Specifies the reduction to apply to the output. Defaults to ``"mean"``. - ``"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. use_softmax: whether to use softmax to transform the original logits into probabilities. If True, softmax is used. If False, sigmoid is used. Defaults to False. Example: >>> import torch >>> from monai.losses import FocalLoss >>> pred = torch.tensor([[1, 0], [0, 1], [1, 0]], dtype=torch.float32) >>> grnd = torch.tensor([[0], [1], [0]], dtype=torch.int64) >>> fl = FocalLoss(to_onehot_y=True) >>> fl(pred, grnd) """ super().__init__(reduction=LossReduction(reduction).value) self.include_background = include_background self.to_onehot_y = to_onehot_y self.gamma = gamma self.alpha = alpha self.weight = weight self.use_softmax = use_softmax weight = torch.as_tensor(weight) if weight is not None else None self.register_buffer("class_weight", weight) self.class_weight: None | torch.Tensor def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: """ Args: input: the shape should be BNH[WD], where N is the number of classes. The input should be the original logits since it will be transformed by a sigmoid/softmax in the forward function. target: the shape should be BNH[WD] or B1H[WD], where N is the number of classes. Raises: ValueError: When input and target (after one hot transform if set) have different shapes. ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. ValueError: When ``self.weight`` is a sequence and the length is not equal to the number of classes. ValueError: When ``self.weight`` is/contains a value that is less than 0. """ n_pred_ch = input.shape[1] if self.to_onehot_y: if n_pred_ch == 1: warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") else: target = one_hot(target, num_classes=n_pred_ch) if not self.include_background: if n_pred_ch == 1: warnings.warn("single channel prediction, `include_background=False` ignored.") else: # if skipping background, removing first channel target = target[:, 1:] input = input[:, 1:] if target.shape != input.shape: raise ValueError(f"ground truth has different shape ({target.shape}) from input ({input.shape})") loss: Optional[torch.Tensor] = None input = input.float() target = target.float() if self.use_softmax: if not self.include_background and self.alpha is not None: self.alpha = None warnings.warn("`include_background=False`, `alpha` ignored when using softmax.") loss = softmax_focal_loss(input, target, self.gamma, self.alpha) else: loss = sigmoid_focal_loss(input, target, self.gamma, self.alpha) num_of_classes = target.shape[1] if self.class_weight is not None and num_of_classes != 1: # make sure the lengths of weights are equal to the number of classes if self.class_weight.ndim == 0: self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes) else: if self.class_weight.shape[0] != num_of_classes: raise ValueError( """the length of the `weight` sequence should be the same as the number of classes. If `include_background=False`, the weight should not include the background category class 0.""" ) if self.class_weight.min() < 0: raise ValueError("the value/values of the `weight` should be no less than 0.") # apply class_weight to loss self.class_weight = self.class_weight.to(loss) broadcast_dims = [-1] + [1] * len(target.shape[2:]) self.class_weight = self.class_weight.view(broadcast_dims) loss = self.class_weight * loss if self.reduction == LossReduction.SUM.value: # Previously there was a mean over the last dimension, which did not # return a compatible BCE loss. To maintain backwards compatible # behavior we have a flag that performs this extra step, disable or # parameterize if necessary. (Or justify why the mean should be there) average_spatial_dims = True if average_spatial_dims: loss = loss.mean(dim=list(range(2, len(target.shape)))) loss = loss.sum() elif self.reduction == LossReduction.MEAN.value: loss = loss.mean() elif self.reduction == LossReduction.NONE.value: pass else: raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].') return loss def softmax_focal_loss( input: torch.Tensor, target: torch.Tensor, gamma: float = 2.0, alpha: Optional[float] = None ) -> torch.Tensor: """ FL(pt) = -alpha * (1 - pt)**gamma * log(pt) where p_i = exp(s_i) / sum_j exp(s_j), t is the target (ground truth) class, and s_j is the unnormalized score for class j. """ input_ls = input.log_softmax(1) loss: torch.Tensor = -(1 - input_ls.exp()).pow(gamma) * input_ls * target if alpha is not None: # (1-alpha) for the background class and alpha for the other classes alpha_fac = torch.tensor([1 - alpha] + [alpha] * (target.shape[1] - 1)).to(loss) broadcast_dims = [-1] + [1] * len(target.shape[2:]) alpha_fac = alpha_fac.view(broadcast_dims) loss = alpha_fac * loss return loss def sigmoid_focal_loss( input: torch.Tensor, target: torch.Tensor, gamma: float = 2.0, alpha: Optional[float] = None ) -> torch.Tensor: """ FL(pt) = -alpha * (1 - pt)**gamma * log(pt) where p = sigmoid(x), pt = p if label is 1 or 1 - p if label is 0 """ # computing binary cross entropy with logits # equivalent to F.binary_cross_entropy_with_logits(input, target, reduction='none') # see also https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/Loss.cpp#L363 loss: torch.Tensor = input - input * target - F.logsigmoid(input) # sigmoid(-i) if t==1; sigmoid(i) if t==0 <=> # 1-sigmoid(i) if t==1; sigmoid(i) if t==0 <=> # 1-p if t==1; p if t==0 <=> # pfac, that is, the term (1 - pt) invprobs = F.logsigmoid(-input * (target * 2 - 1)) # reduced chance of overflow # (pfac.log() * gamma).exp() <=> # pfac.log().exp() ^ gamma <=> # pfac ^ gamma loss = (invprobs * gamma).exp() * loss if alpha is not None: # alpha if t==1; (1-alpha) if t==0 alpha_factor = target * alpha + (1 - target) * (1 - alpha) loss = alpha_factor * loss return loss