sync sum_centerness_targets only when num_gpus > 1

Ranchentx · Aug 26, 2019 · 0145839 · 0145839
1 parent 816bfa6
commit 0145839
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/fcos_core/modeling/rpn/fcos/loss.py b/fcos_core/modeling/rpn/fcos/loss.py
@@ -259,8 +259,13 @@ def __call__(self, locations, box_cls, box_regression, centerness, targets):
 
         if pos_inds.numel() > 0:
             centerness_targets = self.compute_centerness_targets(reg_targets_flatten)
+
             sum_centerness_targets = centerness_targets.sum()
-            sum_centerness_targets = reduce_sum(sum_centerness_targets).item()
+            if num_gpus > 1:
+                sum_centerness_targets = reduce_sum(sum_centerness_targets).item()
+            else:
+                sum_centerness_targets = sum_centerness_targets.item()
+
             reg_loss = self.box_reg_loss_func(
                 box_regression_flatten,
                 reg_targets_flatten,