Introduce DetectionModelHelper.GetLossScale()

Summary: Under a distributed setting (multiple nodes), the loss scale may need to account for the number of nodes participating in the algorithm. This method provides a reasonable way for providing this config. Reviewed By: rbgirshick Differential Revision: D6826976 fbshipit-source-id: d224200710b0f3c3245851874c9f4f3ab871710f
freedomme · Jan 31, 2018 · 3e99b2f · 3e99b2f
1 parent e59c30b
commit 3e99b2f
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 10 deletions.
diff --git a/lib/modeling/FPN.py b/lib/modeling/FPN.py
@@ -412,7 +412,7 @@ def add_fpn_rpn_losses(model):
             'loss_rpn_cls_fpn' + slvl,
             normalize=0,
             scale=(
-                1. / cfg.NUM_GPUS / cfg.TRAIN.RPN_BATCH_SIZE_PER_IM /
+                model.GetLossScale() / cfg.TRAIN.RPN_BATCH_SIZE_PER_IM /
                 cfg.TRAIN.IMS_PER_BATCH
             )
         )
@@ -427,7 +427,7 @@ def add_fpn_rpn_losses(model):
             ],
             'loss_rpn_bbox_fpn' + slvl,
             beta=1. / 9.,
-            scale=1. / cfg.NUM_GPUS
+            scale=model.GetLossScale(),
         )
         loss_gradients.update(
             blob_utils.

diff --git a/lib/modeling/detector.py b/lib/modeling/detector.py
@@ -471,6 +471,13 @@ def _CorrectMomentum(self, correction):
                         scale=correction)
                     workspace.RunOperatorOnce(op)
 
+    def GetLossScale(self):
+        """Allow a way to configure the loss scale dynamically.
+
+        This may be used in a distributed data parallel setting.
+        """
+        return 1.0 / cfg.NUM_GPUS
+
     def AddLosses(self, losses):
         if not isinstance(losses, list):
             losses = [losses]

diff --git a/lib/modeling/fast_rcnn_heads.py b/lib/modeling/fast_rcnn_heads.py
@@ -70,15 +70,15 @@ def add_fast_rcnn_losses(model):
     """Add losses for RoI classification and bounding box regression."""
     cls_prob, loss_cls = model.net.SoftmaxWithLoss(
         ['cls_score', 'labels_int32'], ['cls_prob', 'loss_cls'],
-        scale=1. / cfg.NUM_GPUS
+        scale=model.GetLossScale()
     )
     loss_bbox = model.net.SmoothL1Loss(
         [
             'bbox_pred', 'bbox_targets', 'bbox_inside_weights',
             'bbox_outside_weights'
         ],
         'loss_bbox',
-        scale=1. / cfg.NUM_GPUS
+        scale=model.GetLossScale()
     )
     loss_gradients = blob_utils.get_loss_gradients(model, [loss_cls, loss_bbox])
     model.Accuracy(['cls_prob', 'labels_int32'], 'accuracy_cls')

diff --git a/lib/modeling/mask_rcnn_heads.py b/lib/modeling/mask_rcnn_heads.py
@@ -96,7 +96,7 @@ def add_mask_rcnn_losses(model, blob_mask):
     loss_mask = model.net.SigmoidCrossEntropyLoss(
         [blob_mask, 'masks_int32'],
         'loss_mask',
-        scale=1. / cfg.NUM_GPUS * cfg.MRCNN.WEIGHT_LOSS_MASK
+        scale=model.GetLossScale() * cfg.MRCNN.WEIGHT_LOSS_MASK
     )
     loss_gradients = blob_utils.get_loss_gradients(model, [loss_mask])
     model.AddLosses('loss_mask')

diff --git a/lib/modeling/retinanet_heads.py b/lib/modeling/retinanet_heads.py
@@ -266,7 +266,7 @@ def add_fpn_retinanet_losses(model):
             ],
             'retnet_loss_bbox_' + suffix,
             beta=cfg.RETINANET.BBOX_REG_BETA,
-            scale=1. / cfg.NUM_GPUS * cfg.RETINANET.BBOX_REG_WEIGHT
+            scale=model.GetLossScale() * cfg.RETINANET.BBOX_REG_WEIGHT
         )
         gradients.append(bbox_loss)
         losses.append('retnet_loss_bbox_' + suffix)
@@ -286,7 +286,7 @@ def add_fpn_retinanet_losses(model):
                 ['fl_{}'.format(suffix)],
                 gamma=cfg.RETINANET.LOSS_GAMMA,
                 alpha=cfg.RETINANET.LOSS_ALPHA,
-                scale=(1. / cfg.NUM_GPUS)
+                scale=model.GetLossScale()
             )
             gradients.append(cls_focal_loss)
             losses.append('fl_{}'.format(suffix))
@@ -299,7 +299,7 @@ def add_fpn_retinanet_losses(model):
                 ['fl_{}'.format(suffix), 'retnet_prob_{}'.format(suffix)],
                 gamma=cfg.RETINANET.LOSS_GAMMA,
                 alpha=cfg.RETINANET.LOSS_ALPHA,
-                scale=(1. / cfg.NUM_GPUS),
+                scale=model.GetLossScale(),
             )
             gradients.append(cls_focal_loss)
             losses.append('fl_{}'.format(suffix))

diff --git a/lib/modeling/rpn_heads.py b/lib/modeling/rpn_heads.py
@@ -136,7 +136,7 @@ def add_single_scale_rpn_losses(model):
     loss_rpn_cls = model.net.SigmoidCrossEntropyLoss(
         ['rpn_cls_logits', 'rpn_labels_int32'],
         'loss_rpn_cls',
-        scale=1. / cfg.NUM_GPUS
+        scale=model.GetLossScale()
     )
     loss_rpn_bbox = model.net.SmoothL1Loss(
         [
@@ -145,7 +145,7 @@ def add_single_scale_rpn_losses(model):
         ],
         'loss_rpn_bbox',
         beta=1. / 9.,
-        scale=1. / cfg.NUM_GPUS
+        scale=model.GetLossScale()
     )
     loss_gradients = blob_utils.get_loss_gradients(
         model, [loss_rpn_cls, loss_rpn_bbox]