Apply class weights in maskrcnn loss

PiperOrigin-RevId: 571114756

Author: tensorflower-gardener

official/vision/configs/maskrcnn.py

@@ -223,6 +223,7 @@ class Losses(hyperparams.Config):
   frcnn_class_weight: float = 1.0
   frcnn_box_weight: float = 1.0
   mask_weight: float = 1.0
+  class_weights: Optional[List[float]] = None
 
 
 @dataclasses.dataclass

official/vision/losses/maskrcnn_losses.py

@@ -167,26 +167,40 @@ def __init__(self,
     self._use_binary_cross_entropy = use_binary_cross_entropy
     self._top_k_percent = top_k_percent
 
-  def __call__(self, class_outputs, class_targets):
+  def __call__(self, class_outputs, class_targets, class_weights=None):
     """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.
 
     This function implements the classification loss of the Fast-RCNN.
 
     The classification loss is categorical (or binary) cross entropy on all
     RoIs.
-    Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py  # pylint: disable=line-too-long
+    Reference:
+    https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py
+    # pylint: disable=line-too-long
 
     Args:
       class_outputs: a float tensor representing the class prediction for each
         box with a shape of [batch_size, num_boxes, num_classes].
       class_targets: a float tensor representing the class label for each box
         with a shape of [batch_size, num_boxes].
+      class_weights: A float list containing the weight of each class.
 
     Returns:
       a scalar tensor representing total class loss.
     """
     with tf.name_scope('fast_rcnn_loss'):
+      output_dtype = class_outputs.dtype
       num_classes = class_outputs.get_shape().as_list()[-1]
+      class_weights = (
+          class_weights if class_weights is not None else [1.0] * num_classes
+      )
+      if num_classes != len(class_weights):
+        raise ValueError(
+            'Length of class_weights should be {}'.format(num_classes)
+        )
+
+      class_weights = tf.constant(class_weights, dtype=output_dtype)
+
       class_targets_one_hot = tf.one_hot(
           tf.cast(class_targets, dtype=tf.int32),
           num_classes,
@@ -195,10 +209,15 @@ def __call__(self, class_outputs, class_targets):
         # (batch_size, num_boxes, num_classes)
         cross_entropy_loss = tf.nn.sigmoid_cross_entropy_with_logits(
             labels=class_targets_one_hot, logits=class_outputs)
+        cross_entropy_loss *= class_weights
       else:
         # (batch_size, num_boxes)
         cross_entropy_loss = tf.nn.softmax_cross_entropy_with_logits(
             labels=class_targets_one_hot, logits=class_outputs)
+        class_weight_mask = tf.einsum(
+            '...y,y->...', class_targets_one_hot, class_weights
+        )
+        cross_entropy_loss *= class_weight_mask
 
       if self._top_k_percent < 1.0:
         return self.aggregate_loss_top_k(cross_entropy_loss)

official/vision/losses/maskrcnn_losses_test.py

@@ -77,7 +77,10 @@ def testFastrcnnClassLoss(self, use_binary_cross_entropy):
                                       maxval=num_classes + 1,
                                       dtype=tf.int32)
     loss_fn = maskrcnn_losses.FastrcnnClassLoss(use_binary_cross_entropy)
-    self.assertEqual(tf.rank(loss_fn(class_outputs, class_targets)), 0)
+    class_weights = [1.0] * num_classes
+    self.assertEqual(
+        tf.rank(loss_fn(class_outputs, class_targets, class_weights)), 0
+    )
 
   def testFastrcnnClassLossTopK(self):
     class_targets = tf.constant([[0, 0, 0, 2]])
@@ -87,16 +90,17 @@ def testFastrcnnClassLossTopK(self):
         [100.0, 0.0, 0.0],
         [0.0, 1.0, 0.0],
     ]])
+    class_weights = [1.0, 1.0, 1.0]
     self.assertAllClose(
         maskrcnn_losses.FastrcnnClassLoss(top_k_percent=0.5)(
-            class_outputs, class_targets
+            class_outputs, class_targets, class_weights
         ),
         0.775718,
         atol=1e-4,
     )
     self.assertAllClose(
         maskrcnn_losses.FastrcnnClassLoss(top_k_percent=1.0)(
-            class_outputs, class_targets
+            class_outputs, class_targets, class_weights
         ),
         0.387861,
         atol=1e-4,

official/vision/tasks/maskrcnn.py

@@ -202,8 +202,10 @@ def _build_rpn_losses(
     return rpn_score_loss, rpn_box_loss
 
   def _build_frcnn_losses(
-      self, outputs: Mapping[str, Any],
-      labels: Mapping[str, Any]) -> Tuple[tf.Tensor, tf.Tensor]:
+      self,
+      outputs: Mapping[str, Any],
+      labels: Mapping[str, Any],
+  ) -> Tuple[tf.Tensor, tf.Tensor]:
     """Builds losses for Fast R-CNN."""
     cascade_ious = self.task_config.model.roi_sampler.cascade_iou_thresholds
 
@@ -222,10 +224,19 @@ def _build_frcnn_losses(
     for cas_num in range(num_det_heads):
       frcnn_cls_loss_i = tf.reduce_mean(
           frcnn_cls_loss_fn(
-              outputs['class_outputs_{}'
-                      .format(cas_num) if cas_num else 'class_outputs'],
-              outputs['class_targets_{}'
-                      .format(cas_num) if cas_num else 'class_targets']))
+              outputs[
+                  'class_outputs_{}'.format(cas_num)
+                  if cas_num
+                  else 'class_outputs'
+              ],
+              outputs[
+                  'class_targets_{}'.format(cas_num)
+                  if cas_num
+                  else 'class_targets'
+              ],
+              self.task_config.losses.class_weights,
+          )
+      )
       frcnn_box_loss_i = tf.reduce_mean(
           frcnn_box_loss_fn(
               outputs['box_outputs_{}'.format(cas_num
@@ -257,27 +268,28 @@ def build_losses(self,
                    labels: Mapping[str, Any],
                    aux_losses: Optional[Any] = None) -> Dict[str, tf.Tensor]:
     """Builds Mask R-CNN losses."""
+    loss_params = self.task_config.losses
     rpn_score_loss, rpn_box_loss = self._build_rpn_losses(outputs, labels)
     frcnn_cls_loss, frcnn_box_loss = self._build_frcnn_losses(outputs, labels)
     if self.task_config.model.include_mask:
       mask_loss = self._build_mask_loss(outputs)
     else:
       mask_loss = tf.constant(0.0, dtype=tf.float32)
 
-    params = self.task_config
     model_loss = (
-        params.losses.rpn_score_weight * rpn_score_loss +
-        params.losses.rpn_box_weight * rpn_box_loss +
-        params.losses.frcnn_class_weight * frcnn_cls_loss +
-        params.losses.frcnn_box_weight * frcnn_box_loss +
-        params.losses.mask_weight * mask_loss)
+        loss_params.rpn_score_weight * rpn_score_loss
+        + loss_params.rpn_box_weight * rpn_box_loss
+        + loss_params.frcnn_class_weight * frcnn_cls_loss
+        + loss_params.frcnn_box_weight * frcnn_box_loss
+        + loss_params.mask_weight * mask_loss
+    )
 
     total_loss = model_loss
     if aux_losses:
       reg_loss = tf.reduce_sum(aux_losses)
       total_loss = model_loss + reg_loss
 
-    total_loss = params.losses.loss_weight * total_loss
+    total_loss = loss_params.loss_weight * total_loss
     losses = {
         'total_loss': total_loss,
         'rpn_score_loss': rpn_score_loss,