feat(losses): Add regularizers to losses

Author: okunator

cellseg_models_pytorch/losses/criterions/ce.py

@@ -1,31 +1,47 @@
 import torch
-import torch.nn as nn
+import torch.nn.functional as F
 
+from ...utils import tensor_one_hot
 from ..weighted_base_loss import WeightedBaseLoss
 
+__all__ = ["CELoss"]
+
 
 class CELoss(WeightedBaseLoss):
     def __init__(
-        self, edge_weight: float = None, class_weights: torch.Tensor = None, **kwargs
+        self,
+        apply_sd: bool = False,
+        apply_ls: bool = False,
+        apply_svls: bool = False,
+        edge_weight: float = None,
+        class_weights: torch.Tensor = None,
+        **kwargs,
     ) -> None:
         """Cross-Entropy loss with weighting.
 
         Parameters
         ----------
-            edge_weight : float, default=none
+            apply_sd : bool, default=False
+                If True, Spectral decoupling regularization will be applied  to the
+                loss matrix.
+            apply_ls : bool, default=False
+                If True, Label smoothing will be applied to the target.
+            apply_svls : bool, default=False
+                If True, spatially varying label smoothing will be applied to the target
+            edge_weight : float, default=None
                 Weight that is added to object borders.
             class_weights : torch.Tensor, default=None
                 Class weights. A tensor of shape (n_classes,).
         """
-        super().__init__(class_weights, edge_weight)
-        self.loss = nn.CrossEntropyLoss(reduction="none", weight=class_weights)
+        super().__init__(apply_sd, apply_ls, apply_svls, class_weights, edge_weight)
+        self.eps = 1e-8
 
     def forward(
         self,
         yhat: torch.Tensor,
         target: torch.Tensor,
         target_weight: torch.Tensor = None,
-        **kwargs
+        **kwargs,
     ) -> torch.Tensor:
         """Compute the cross entropy loss.
 
@@ -43,7 +59,28 @@ def forward(
             torch.Tensor:
                 Computed CE loss (scalar).
         """
-        loss = self.loss(yhat, target)  # (B, H, W)
+        input_soft = F.softmax(yhat, dim=1) + self.eps  # (B, C, H, W)
+        num_classes = yhat.shape[1]
+        target_one_hot = tensor_one_hot(target, num_classes)  # (B, C, H, W)
+        assert target_one_hot.shape == yhat.shape
+
+        if self.apply_svls:
+            target_one_hot = self.apply_svls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
+        if self.apply_ls:
+            target_one_hot = self.apply_ls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
+        loss = -torch.sum(target_one_hot * torch.log(input_soft), dim=1)  # (B, H, W)
+
+        if self.apply_sd:
+            loss = self.apply_spectral_decouple(loss, yhat)
+
+        if self.class_weights is not None:
+            loss = self.apply_class_weights(loss, target)
 
         if self.edge_weight is not None:
             loss = self.apply_edge_weights(loss, target_weight)

cellseg_models_pytorch/losses/criterions/dice.py

@@ -5,31 +5,46 @@
 
 from ..weighted_base_loss import WeightedBaseLoss
 
+__all__ = ["DiceLoss"]
+
 
 class DiceLoss(WeightedBaseLoss):
     def __init__(
-        self, edge_weight: float = None, class_weights: torch.Tensor = None, **kwargs
+        self,
+        apply_sd: bool = False,
+        apply_ls: bool = False,
+        apply_svls: bool = False,
+        edge_weight: float = None,
+        class_weights: torch.Tensor = None,
+        **kwargs,
     ) -> None:
         """Sørensen-Dice Coefficient Loss.
 
         Optionally applies weights at the object edges and classes.
 
         Parameters
         ----------
+            apply_sd : bool, default=False
+                If True, Spectral decoupling regularization will be applied  to the
+                loss matrix.
+            apply_ls : bool, default=False
+                If True, Label smoothing will be applied to the target.
+            apply_svls : bool, default=False
+                If True, spatially varying label smoothing will be applied to the target
             edge_weight : float, default=none
                 Weight that is added to object borders.
             class_weights : torch.Tensor, default=None
                 Class weights. A tensor of shape (n_classes,).
         """
-        super().__init__(class_weights, edge_weight)
-        self.eps = 1e-6
+        super().__init__(apply_sd, apply_ls, apply_svls, class_weights, edge_weight)
+        self.eps = 1e-8
 
     def forward(
         self,
         yhat: torch.Tensor,
         target: torch.Tensor,
         target_weight: torch.Tensor = None,
-        **kwargs
+        **kwargs,
     ) -> torch.Tensor:
         """Compute the DICE coefficient.
 
@@ -48,12 +63,26 @@ def forward(
                 Computed DICE loss (scalar).
         """
         yhat_soft = F.softmax(yhat, dim=1)
-        target_one_hot = tensor_one_hot(target, n_classes=yhat.shape[1])
+        num_classes = yhat.shape[1]
+        target_one_hot = tensor_one_hot(target, n_classes=num_classes)
         assert target_one_hot.shape == yhat.shape
 
+        if self.apply_svls:
+            target_one_hot = self.apply_svls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
+        if self.apply_ls:
+            target_one_hot = self.apply_ls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
         intersection = torch.sum(yhat_soft * target_one_hot, 1)
         union = torch.sum(yhat_soft + target_one_hot, 1)
-        dice = 2.0 * intersection / union.clamp_min(self.eps)
+        dice = 2.0 * intersection / union.clamp_min(self.eps)  # (B, H, W)
+
+        if self.apply_sd:
+            dice = self.apply_spectral_decouple(dice, yhat)
 
         if self.class_weights is not None:
             dice = self.apply_class_weights(dice, target)

cellseg_models_pytorch/losses/criterions/focal.py

@@ -1,8 +1,7 @@
 import torch
 import torch.nn.functional as F
 
-from cellseg_models_pytorch.utils import tensor_one_hot
-
+from ...utils import tensor_one_hot
 from ..weighted_base_loss import WeightedBaseLoss
 
 
@@ -11,6 +10,9 @@ def __init__(
         self,
         alpha: float = 0.5,
         gamma: float = 2.0,
+        apply_sd: bool = False,
+        apply_ls: bool = False,
+        apply_svls: bool = False,
         edge_weight: float = None,
         class_weights: torch.Tensor = None,
         **kwargs
@@ -19,23 +21,31 @@ def __init__(
 
         https://arxiv.org/abs/1708.02002
 
-        Optionally applies weights at the object edges and classes.
+        Optionally applies, label smoothing, spatially varying label smoothing or
+        weights at the object edges or class weights to the loss.
 
         Parameters
         ----------
             alpha : float, default=0.5
                 Weight factor b/w [0,1].
             gamma : float, default=2.0
                 Focusing factor.
+            apply_sd : bool, default=False
+                If True, Spectral decoupling regularization will be applied  to the
+                loss matrix.
+            apply_ls : bool, default=False
+                If True, Label smoothing will be applied to the target.
+            apply_svls : bool, default=False
+                If True, spatially varying label smoothing will be applied to the target
             edge_weight : float, default=none
                 Weight that is added to object borders.
             class_weights : torch.Tensor, default=None
                 Class weights. A tensor of shape (n_classes,).
         """
-        super().__init__(class_weights, edge_weight)
+        super().__init__(apply_sd, apply_ls, apply_svls, class_weights, edge_weight)
         self.alpha = alpha
         self.gamma = gamma
-        self.eps = 1e-6
+        self.eps = 1e-8
 
     def forward(
         self,
@@ -65,12 +75,25 @@ def forward(
         target_one_hot = tensor_one_hot(target, num_classes)  # (B, C, H, W)
         assert target_one_hot.shape == yhat.shape
 
+        if self.apply_svls:
+            target_one_hot = self.apply_svls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
+        if self.apply_ls:
+            target_one_hot = self.apply_ls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
         weight = (1.0 - input_soft) ** self.gamma
         focal = self.alpha * weight * torch.log(input_soft)
         focal = target_one_hot * focal
 
         loss = -torch.sum(focal, dim=1)  # to (B, H, W)
 
+        if self.apply_sd:
+            loss = self.apply_spectral_decouple(loss, yhat)
+
         if self.class_weights is not None:
             loss = self.apply_class_weights(loss, target)
 

cellseg_models_pytorch/losses/criterions/iou.py

@@ -8,28 +8,41 @@
 
 class IoULoss(WeightedBaseLoss):
     def __init__(
-        self, edge_weight: float = None, class_weights: torch.Tensor = None, **kwargs
+        self,
+        apply_sd: bool = False,
+        apply_ls: bool = False,
+        apply_svls: bool = False,
+        edge_weight: float = None,
+        class_weights: torch.Tensor = None,
+        **kwargs,
     ) -> None:
         """Intersection over union loss.
 
         Optionally applies weights at the object edges and classes.
 
         Parameters
         ----------
+            apply_sd : bool, default=False
+                If True, Spectral decoupling regularization will be applied  to the
+                loss matrix.
+            apply_ls : bool, default=False
+                If True, Label smoothing will be applied to the target.
+            apply_svls : bool, default=False
+                If True, spatially varying label smoothing will be applied to the target
             edge_weight : float, default=none
                 Weight that is added to object borders.
             class_weights : torch.Tensor, default=None
                 Class weights. A tensor of shape (n_classes,).
         """
-        super().__init__(class_weights, edge_weight)
-        self.eps = 1e-6
+        super().__init__(apply_sd, apply_ls, apply_svls, class_weights, edge_weight)
+        self.eps = 1e-8
 
     def forward(
         self,
         yhat: torch.Tensor,
         target: torch.Tensor,
         target_weight: torch.Tensor = None,
-        **kwargs
+        **kwargs,
     ) -> torch.Tensor:
         """Compute the IoU loss.
 
@@ -48,13 +61,27 @@ def forward(
                 Computed IoU loss (scalar).
         """
         yhat_soft = F.softmax(yhat, dim=1)
-        target_one_hot = tensor_one_hot(target, n_classes=yhat.shape[1])
+        num_classes = yhat.shape[1]
+        target_one_hot = tensor_one_hot(target, n_classes=num_classes)
         assert target_one_hot.shape == yhat.shape
 
+        if self.apply_svls:
+            target_one_hot = self.apply_svls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
+        if self.apply_ls:
+            target_one_hot = self.apply_ls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
         intersection = torch.sum(yhat_soft * target_one_hot, 1)  # to (B, H, W)
         union = torch.sum(yhat_soft + target_one_hot, 1)  # to (B, H, W)
         iou = intersection / union.clamp_min(self.eps)
 
+        if self.apply_sd:
+            iou = self.apply_spectral_decouple(iou, yhat)
+
         if self.class_weights is not None:
             iou = self.apply_class_weights(iou, target)
 

cellseg_models_pytorch/losses/criterions/mse.py

@@ -8,25 +8,38 @@
 
 class MSE(WeightedBaseLoss):
     def __init__(
-        self, edge_weight: float = None, class_weights: torch.Tensor = None, **kwargs
+        self,
+        apply_sd: bool = False,
+        apply_ls: bool = False,
+        apply_svls: bool = False,
+        edge_weight: float = None,
+        class_weights: torch.Tensor = None,
+        **kwargs,
     ) -> None:
         """MSE-loss.
 
         Parameters
         ----------
+            apply_sd : bool, default=False
+                If True, Spectral decoupling regularization will be applied  to the
+                loss matrix.
+            apply_ls : bool, default=False
+                If True, Label smoothing will be applied to the target.
+            apply_svls : bool, default=False
+                If True, spatially varying label smoothing will be applied to the target
             edge_weight : float, default=none
                 Weight that is added to object borders.
             class_weights : torch.Tensor, default=None
                 Class weights. A tensor of shape (n_classes,).
         """
-        super().__init__(class_weights, edge_weight)
+        super().__init__(apply_sd, apply_ls, apply_svls, class_weights, edge_weight)
 
     def forward(
         self,
         yhat: torch.Tensor,
         target: torch.Tensor,
         target_weight: torch.Tensor = None,
-        **kwargs
+        **kwargs,
     ) -> torch.Tensor:
         """Compute the MSE-loss.
 
@@ -45,15 +58,30 @@ def forward(
                 Computed MSE loss (scalar).
         """
         target_one_hot = target
+        num_classes = yhat.shape[1]
+
         if target.size() != yhat.size():
             if target.dtype == torch.float32:
                 target_one_hot = target.unsqueeze(1)
             else:
-                target_one_hot = tensor_one_hot(target, yhat.shape[1])
+                target_one_hot = tensor_one_hot(target, num_classes)
+
+        if self.apply_svls:
+            target_one_hot = self.apply_svls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
+
+        if self.apply_ls:
+            target_one_hot = self.apply_ls_to_target(
+                target_one_hot, num_classes, **kwargs
+            )
 
         mse = F.mse_loss(yhat, target_one_hot, reduction="none")  # (B, C, H, W)
         mse = torch.mean(mse, dim=1)  # to (B, H, W)
 
+        if self.apply_sd:
+            mse = self.apply_spectral_decouple(mse, yhat)
+
         if self.class_weights is not None:
             mse = self.apply_class_weights(mse, target)