stdc bugfix for 1.7.5

Author: shaydeci

src/super_gradients/training/exceptions/loss_exceptions.py

@@ -0,0 +1,24 @@
+class IllegalRangeForLossAttributeException(Exception):
+    """
+    Exception raised illegal value (i.e not in range) for _Loss attribute.
+
+    Attributes:
+        message -- explanation of the error
+    """
+
+    def __init__(self, range_vals: tuple, attr_name: str):
+        self.message = attr_name + " must be in range " + str(range_vals)
+        super().__init__(self.message)
+
+
+class RequiredLossComponentReductionException(Exception):
+    """
+    Exception raised illegal reduction for _Loss component.
+
+    Attributes:
+        message -- explanation of the error
+    """
+
+    def __init__(self, component_name: str, reduction: str, required_reduction: str):
+        self.message = component_name + ".reduction must be " + required_reduction + ", got" + reduction
+        super().__init__(self.message)

src/super_gradients/training/losses/ohem_ce_loss.py

@@ -1,32 +1,44 @@
 import torch
 from torch import nn
 from torch.nn.modules.loss import _Loss
+from super_gradients.training.exceptions.loss_exceptions import IllegalRangeForLossAttributeException, RequiredLossComponentReductionException
 
 
-class OhemCELoss(_Loss):
+class OhemLoss(_Loss):
     """
-    OhemCELoss - Online Hard Example Mining Cross Entropy Loss
+    OhemLoss - Online Hard Example Mining Cross Entropy Loss
     """
+
     def __init__(self,
                  threshold: float,
                  mining_percent: float = 0.1,
                  ignore_lb: int = -100,
-                 num_pixels_exclude_ignored: bool = True):
+                 num_pixels_exclude_ignored: bool = True,
+                 criteria: _Loss = None):
         """
         :param threshold: Sample below probability threshold, is considered hard.
         :param num_pixels_exclude_ignored: How to calculate total pixels from which extract mining percent of the
          samples.
+        :param ignore_lb: label index to be ignored in loss calculation.
+        :param criteria: loss to mine the examples from.
+
          i.e for num_pixels=100, ignore_pixels=30, mining_percent=0.1:
          num_pixels_exclude_ignored=False => num_mining = 100 * 0.1 = 10
          num_pixels_exclude_ignored=True  => num_mining = (100 - 30) * 0.1 = 7
         """
         super().__init__()
-        assert 0 <= mining_percent <= 1, "mining percent should be a value from 0 to 1"
+
+        if mining_percent < 0 or mining_percent > 1:
+            raise IllegalRangeForLossAttributeException((0, 1), "mining percent")
+
         self.thresh = -torch.log(torch.tensor(threshold, dtype=torch.float))
         self.mining_percent = mining_percent
-        self.ignore_lb = -100 if ignore_lb is None or ignore_lb < 0 else ignore_lb
+        self.ignore_lb = ignore_lb
         self.num_pixels_exclude_ignored = num_pixels_exclude_ignored
-        self.criteria = nn.CrossEntropyLoss(ignore_index=self.ignore_lb, reduction='none')
+
+        if criteria.reduction != 'none':
+            raise RequiredLossComponentReductionException("criteria", criteria.reduction, 'none')
+        self.criteria = criteria
 
     def forward(self, logits, labels):
         loss = self.criteria(logits, labels).view(-1)
@@ -52,3 +64,46 @@ def forward(self, logits, labels):
         else:
             loss = loss[:num_mining]
         return torch.mean(loss)
+
+
+class OhemCELoss(OhemLoss):
+    """
+    OhemLoss - Online Hard Example Mining Cross Entropy Loss
+    """
+
+    def __init__(self,
+                 threshold: float,
+                 mining_percent: float = 0.1,
+                 ignore_lb: int = -100,
+                 num_pixels_exclude_ignored: bool = True):
+        ignore_lb = -100 if ignore_lb is None or ignore_lb < 0 else ignore_lb
+        criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='none')
+        super(OhemCELoss, self).__init__(threshold=threshold,
+                                         mining_percent=mining_percent,
+                                         ignore_lb=ignore_lb,
+                                         num_pixels_exclude_ignored=num_pixels_exclude_ignored,
+                                         criteria=criteria)
+
+
+class OhemBCELoss(OhemLoss):
+    """
+    OhemBCELoss - Online Hard Example Mining Binary Cross Entropy Loss
+    """
+
+    def __init__(self,
+                 threshold: float,
+                 mining_percent: float = 0.1,
+                 ignore_lb: int = -100,
+                 num_pixels_exclude_ignored: bool = True, ):
+        super(OhemBCELoss, self).__init__(threshold=threshold,
+                                          mining_percent=mining_percent,
+                                          ignore_lb=ignore_lb,
+                                          num_pixels_exclude_ignored=num_pixels_exclude_ignored,
+                                          criteria=nn.BCEWithLogitsLoss(reduction='none'))
+
+    def forward(self, logits, labels):
+
+        # REMOVE SINGLE CLASS CHANNEL WHEN DEALING WITH BINARY DATA
+        if logits.shape[1] == 1:
+            logits = logits.squeeze(1)
+        return super(OhemBCELoss, self).forward(logits, labels.float())

src/super_gradients/training/losses/stdc_loss.py

@@ -3,7 +3,7 @@
 import torch.nn.functional as F
 from super_gradients.training.utils.segmentation_utils import to_one_hot
 from torch.nn.modules.loss import _Loss
-from super_gradients.training.losses.ohem_ce_loss import OhemCELoss
+from super_gradients.training.losses.ohem_ce_loss import OhemCELoss, OhemBCELoss, OhemLoss
 from super_gradients.training.losses.dice_loss import BinaryDiceLoss
 from typing import Union, Tuple
 
@@ -46,8 +46,10 @@ def __init__(self,
     def forward(self, gt_masks: torch.Tensor):
         if self.device is None:
             self._set_kernels_to_device(gt_masks.device)
-
-        one_hot = to_one_hot(gt_masks, self.num_classes, self.ignore_label).float()
+        if self.num_classes > 1:
+            one_hot = to_one_hot(gt_masks, self.num_classes, self.ignore_label).float()
+        else:
+            one_hot = gt_masks.unsqueeze(1).float()
         # create binary detail maps using filters withs strides of 1, 2 and 4.
         boundary_targets = F.conv2d(one_hot, self.laplacian_kernel, stride=1, padding=1, groups=self.num_classes)
         boundary_targets_x2 = F.conv2d(one_hot, self.laplacian_kernel, stride=2, padding=1, groups=self.num_classes)
@@ -123,7 +125,8 @@ def __init__(self,
                  mining_percent: float = 0.1,
                  detail_threshold: float = 1.,
                  learnable_fusing_kernel: bool = True,
-                 ignore_index: int = None):
+                 ignore_index: int = None,
+                 ohem_criteria: OhemLoss = None):
         """
         :param threshold: Online hard-mining probability threshold.
         :param num_aux_heads: num of auxiliary heads.
@@ -133,6 +136,8 @@ def __init__(self,
         :param mining_percent: mining percentage.
         :param detail_threshold: detail threshold to create binary details features in DetailLoss.
         :param learnable_fusing_kernel: whether DetailAggregateModule params are learnable or not.
+        :param ohem_criteria: OhemLoss criterion component of STDC. When none is given, it will be derrived according
+         to num_classes (i.e OhemCELoss if num_classes > 1 and OhemBCELoss otherwise).
         """
         super().__init__()
 
@@ -151,7 +156,11 @@ def __init__(self,
                                                        learnable_fusing_kernel=learnable_fusing_kernel)
             self.detail_loss = DetailLoss(weights=detail_weights)
 
-        self.ce_ohem = OhemCELoss(threshold=threshold, mining_percent=mining_percent, ignore_lb=ignore_index)
+        if ohem_criteria is None:
+            ohem_criteria = OhemCELoss(threshold=threshold, mining_percent=mining_percent, ignore_lb=ignore_index) if num_classes > 1 else OhemBCELoss(threshold=threshold, mining_percent=mining_percent)
+
+        self.ce_ohem = ohem_criteria
+        self.num_classes = num_classes
 
     def forward(self, preds: Tuple[torch.Tensor], target: torch.Tensor):
         """

src/super_gradients/training/metrics/segmentation_metrics.py

@@ -119,6 +119,9 @@ def __init__(self, dist_sync_on_step=True, ignore_index=None):
         self.component_names = ["target_IOU", "background_IOU", "mean_IOU"]
 
     def update(self, preds, target: torch.Tensor):
+        # WHEN DEALING WITH MULTIPLE OUTPUTS- OUTPUTS[0] IS THE MAIN SEGMENTATION MAP
+        if isinstance(preds, tuple):
+            preds = preds[0]
         super().update(preds=torch.sigmoid(preds), target=target.long())
 
     def compute(self):

src/super_gradients/training/utils/callbacks.py

@@ -494,7 +494,10 @@ def __init__(self, phase: Phase, freq: int, batch_idx: int = 0, last_img_idx_in_
 
     def __call__(self, context: PhaseContext):
         if context.epoch % self.freq == 0 and context.batch_idx == self.batch_idx:
-            preds = context.preds.clone()
+            if isinstance(context.preds, tuple):
+                preds = context.preds[0].clone()
+            else:
+                preds = context.preds.clone()
             batch_imgs = BinarySegmentationVisualization.visualize_batch(context.inputs, preds, context.target, self.batch_idx)
             batch_imgs = [cv2.cvtColor(image, cv2.COLOR_BGR2RGB) for image in batch_imgs]
             batch_imgs = np.stack(batch_imgs)