flop counter: clean-up of tracing logic

Summary:
1. remove check of _get_trace_graph: the old version was removed in pytorch for long enough
2.  add a test case about dataparallel
3. move _named_module_without_dup to a separate function so the tracing function looks simpler

Reviewed By: ericmintun

Differential Revision: D32227000

fbshipit-source-id: fada03028b48db33ef2e2058853c6dda39750251

Author: ppwwyyxx

fvcore/nn/jit_analysis.py

@@ -8,7 +8,7 @@
 from copy import copy
 from dataclasses import dataclass
 from numbers import Number
-from typing import Any, Dict, List, Optional, Set, Tuple, TypeVar, Union
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, TypeVar, Union
 
 import numpy as np
 import torch
@@ -99,6 +99,18 @@ class Statistics:
     uncalled_mods: "Set[str]"
 
 
+def _named_modules_without_dup(model: nn.Module) -> Iterator[Tuple[str, nn.Module]]:
+    """
+    Like .named_modules(), but the results are slightly different for
+    some wrapped models.
+    """
+    seen = set()
+    for name, mod in _named_modules_with_dup(model):
+        if mod not in seen:
+            seen.add(mod)
+            yield name, mod
+
+
 def _get_scoped_trace_graph(
     module: nn.Module,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
@@ -139,7 +151,6 @@ def __call__(self, module: nn.Module, inputs: Any, outputs: Any) -> Any:
                 tracing_state.pop_scope()
             return outputs
 
-    seen = set()
     hook_handles: List[Any] = []
 
     def register_hooks(mod: nn.Module, name: str) -> None:
@@ -149,31 +160,21 @@ def register_hooks(mod: nn.Module, name: str) -> None:
         hook_handles.append(prehook)
         hook_handles.append(posthook)
 
-    # Torch script does not support parallel torch models, but we still
-    # want the scope names to be correct for the complete module.
+    # Unwrap DDP, but correct the scope names for the root module.
     if isinstance(
         module, (nn.parallel.distributed.DistributedDataParallel, nn.DataParallel)
     ):
-
         # Since DataParallel just wraps the model, add an extra set of hooks
         # to the model it wraps to account for the wrapper. Then trace it.
         root_name = aliases[module]
         module = module.module
         register_hooks(module, root_name)
 
-    # We don't need the duplication here, but self._model.named_modules()
-    # gives slightly different results for some wrapped models.
-    for name, mod in _named_modules_with_dup(module):
-        if mod not in seen:
-            name = aliases[mod]
-            register_hooks(mod, name)
-            seen.add(mod)
+    for name, mod in _named_modules_without_dup(module):
+        name = aliases[mod]
+        register_hooks(mod, name)
 
-    if hasattr(torch.jit, "get_trace_graph"):
-        trace, _ = torch.jit.get_trace_graph(module, inputs)
-        graph = trace.graph()
-    else:
-        graph, _ = _get_trace_graph(module, inputs)
+    graph, _ = _get_trace_graph(module, inputs)
 
     for handle in hook_handles:
         handle.remove()

tests/test_jit_model_analysis.py

@@ -566,6 +566,14 @@ def test_data_parallel(self) -> None:
         # Test no uncalled modules
         self.assertEqual(analyzer.uncalled_modules(), set())
 
+    def test_data_parallel_root_scope(self) -> None:
+        # A test case discussed in D32227000
+        model = nn.DataParallel(nn.Linear(10, 10))
+        for mode in ["caller", "owner"]:
+            flop = FlopCountAnalysis(model, (torch.randn(10, 10),))
+            flop.ancestor_mode(mode)
+            self.assertEqual(flop.total(), 1000)
+
     def test_unsupported_ops(self) -> None:
         """
         Tests per-module recording of unsupported operations.