Add find_unused_parameters

Summary: can be called when DDP fails to check what params cause the failure

Differential Revision: D31463117

fbshipit-source-id: c12b6c82f16916d4ff9175f2d93d951d4f313350

Author: ppwwyyxx

detectron2/utils/analysis.py

@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 
 import typing
+from typing import Any, List
 import fvcore
 from fvcore.nn import activation_count, flop_count, parameter_count, parameter_count_table
 from torch import nn
@@ -151,3 +152,36 @@ def _wrapper_count_operators(
         ret = ret[0]
     model.train(old_train)
     return ret
+
+
+def find_unused_parameters(model: nn.Module, inputs: Any) -> List[str]:
+    """
+    Given a model, find parameters that do not contribute
+    to the loss.
+
+    Args:
+        model: a model in training mode that returns losses
+        inputs: argument or a tuple of arguments. Inputs of the model
+
+    Returns:
+        list[str]: the name of unused parameters
+    """
+    assert model.training
+    for _, prm in model.named_parameters():
+        prm.grad = None
+
+    if isinstance(inputs, tuple):
+        losses = model(*inputs)
+    else:
+        losses = model(inputs)
+
+    if isinstance(losses, dict):
+        losses = sum(losses.values())
+    losses.backward()
+
+    unused: List[str] = []
+    for name, prm in model.named_parameters():
+        if prm.grad is None:
+            unused.append(name)
+        prm.grad = None
+    return unused

tests/test_model_analysis.py

@@ -3,8 +3,9 @@
 
 import unittest
 import torch
+from torch import nn
 
-from detectron2.utils.analysis import flop_count_operators, parameter_count
+from detectron2.utils.analysis import find_unused_parameters, flop_count_operators, parameter_count
 from detectron2.utils.testing import get_model_no_weights
 
 
@@ -42,3 +43,19 @@ def test_param_count(self):
         res = parameter_count(self.model)
         self.assertTrue(res[""], 41699936)
         self.assertTrue(res["backbone"], 26799296)
+
+
+class UnusedParamTest(unittest.TestCase):
+    def test_unused(self):
+        class TestMod(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fc1 = nn.Linear(10, 10)
+                self.t = nn.Linear(10, 10)
+
+            def forward(self, x):
+                return self.fc1(x).mean()
+
+        m = TestMod()
+        ret = find_unused_parameters(m, torch.randn(10, 10))
+        self.assertEqual(set(ret), {"t.weight", "t.bias"})