fix(rtdetrv2): Fix incorrect method to inspect total norm during training stage (#9)

Author: guyleaf

rtdetrv2_pytorch/rtdetrv2/misc/dist_utils.py

@@ -228,6 +228,18 @@ def reduce_dict(data, avg=True):
         return {k: v for k, v in zip(keys, values)}
 
 
+def all_reduce(data, op=tdist.ReduceOp.SUM):
+    """
+    Run all_reduce on torch.Tensor data
+    Args:
+        data: torch.Tensor
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return
+    tdist.all_reduce(data, op=op)
+
+
 def all_gather(data):
     """
     Run all_gather on arbitrary picklable data (not necessarily tensors)

rtdetrv2_pytorch/rtdetrv2/solver/det_engine.py

@@ -61,12 +61,13 @@ def train_one_epoch(
             loss = sum(loss_dict.values())
             scaler.scale(loss).backward()
 
+            scaler.unscale_(optimizer)
             if max_norm > 0:
-                scaler.unscale_(optimizer)
                 total_norm = torch.nn.utils.clip_grad_norm_(
                     model.parameters(), max_norm
                 )
-                loss_dict["grad_norm"] = total_norm
+            else:
+                total_norm = torch.nn.utils.get_total_norm(model.parameters())
 
             scaler.step(optimizer)
             scaler.update()
@@ -84,7 +85,8 @@ def train_one_epoch(
                 total_norm = torch.nn.utils.clip_grad_norm_(
                     model.parameters(), max_norm
                 )
-                loss_dict["grad_norm"] = total_norm
+            else:
+                total_norm = torch.nn.utils.get_total_norm(model.parameters())
 
             optimizer.step()
 
@@ -103,6 +105,10 @@ def train_one_epoch(
             print(loss_dict_reduced)
             sys.exit(1)
 
+        # collect other values for logging
+        dist_utils.all_reduce(total_norm)
+        loss_dict_reduced["grad_norm"] = total_norm
+
         metric_logger.update(loss=loss_value, **loss_dict_reduced)
         metric_logger.update(lr=optimizer.param_groups[0]["lr"])