[Fix] Fix tracer

mmrazor/structures/tracer/loss_calculator/image_classifier_loss_calculator.py

@@ -13,9 +13,17 @@
 @TASK_UTILS.register_module()
 class ImageClassifierPseudoLoss:
     """Calculate the pseudo loss to trace the topology of a `ImageClassifier`
-    in MMClassification with `BackwardTracer`."""
+    in MMClassification with `BackwardTracer`.
+
+    Args:
+        input_shape (Tuple): The shape of the pseudo input. Defaults to
+            (2, 3, 224, 224).
+    """
+
+    def __init__(self, input_shape=(2, 3, 224, 224)):
+        self.input_shape = input_shape
 
     def __call__(self, model: ImageClassifier) -> torch.Tensor:
-        pseudo_img = torch.rand(1, 3, 224, 224)
+        pseudo_img = torch.rand(self.input_shape)
         pseudo_output = model(pseudo_img)
-        return sum(pseudo_output)
+        return pseudo_output.sum()

mmrazor/structures/tracer/loss_calculator/single_stage_detector_loss_calculator.py

@@ -1,17 +1,31 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
-from mmdet.models import BaseDetector
 
 from mmrazor.registry import TASK_UTILS
 
+try:
+    from mmdet.models import SingleStageDetector
+except ImportError:
+    from mmrazor.utils import get_placeholder
+    SingleStageDetector = get_placeholder('mmdet')
+
 
-# todo: adapt to mmdet 2.0
 @TASK_UTILS.register_module()
 class SingleStageDetectorPseudoLoss:
+    """Calculate the pseudo loss to trace the topology of a
+    `SingleStageDetector` in MMDetection with `BackwardTracer`.
+
+    Args:
+        input_shape (Tuple): The shape of the pseudo input. Defaults to
+            (2, 3, 224, 224).
+    """
+
+    def __init__(self, input_shape=(2, 3, 224, 224)):
+        self.input_shape = input_shape
 
-    def __call__(self, model: BaseDetector) -> torch.Tensor:
-        pseudo_img = torch.rand(1, 3, 224, 224)
-        pseudo_output = model.forward_dummy(pseudo_img)
+    def __call__(self, model: SingleStageDetector) -> torch.Tensor:
+        pseudo_img = torch.rand(self.input_shape)
+        pseudo_output = model(pseudo_img)
         out = torch.tensor(0.)
         for levels in pseudo_output:
             out += sum([level.sum() for level in levels])

mmrazor/structures/tracer/parsers.py

@@ -118,27 +118,16 @@ def parse_cat(tracer, grad_fn, module2name, param2module, cur_path,
         >>> # ``out`` is obtained by concatenating two tensors
     """
     parents = grad_fn.next_functions
-    concat_id = '_'.join([str(id(p)) for p in parents])
-    name = f'concat_{concat_id}'
-    # If a module is not a shared module and it has been visited during
-    # forward, its parent modules must have been traced already.
-    # However, a shared module will be visited more than once during
-    # forward, so it is still need to be traced even if it has been
-    # visited.
-    if (name in visited and visited[name] and name not in shared_module):
-        pass
-    else:
-        visited[name] = True
-        sub_path_lists = list()
-        for i, parent in enumerate(parents):
-            sub_path_list = PathList()
-            tracer.backward_trace(parent, module2name, param2module, Path(),
-                                  sub_path_list, visited, shared_module)
-            sub_path_lists.append(sub_path_list)
-        cur_path.append(PathConcatNode(name, sub_path_lists))
-
-        result_paths.append(copy.deepcopy(cur_path))
-        cur_path.pop(-1)
+    sub_path_lists = list()
+    for i, parent in enumerate(parents):
+        sub_path_list = PathList()
+        tracer.backward_trace(parent, module2name, param2module, Path(),
+                              sub_path_list, visited, shared_module)
+        sub_path_lists.append(sub_path_list)
+    cur_path.append(PathConcatNode('CatNode', sub_path_lists))
+
+    result_paths.append(copy.deepcopy(cur_path))
+    cur_path.pop(-1)
 
 
 def parse_norm(tracer, grad_fn, module2name, param2module, cur_path,
@@ -174,6 +163,8 @@ def parse_norm(tracer, grad_fn, module2name, param2module, cur_path,
 
 
 DEFAULT_BACKWARD_TRACER: Dict[str, Callable] = {
+    'ConvolutionBackward': parse_conv,
+    'SlowConv2DBackward': parse_conv,
     'ThnnConv2DBackward': parse_conv,
     'CudnnConvolutionBackward': parse_conv,
     'MkldnnConvolutionBackward': parse_conv,

tests/test_core/test_tracer/test_backward_tracer.py

@@ -56,6 +56,31 @@ def forward(self, x: Tensor) -> Tensor:
         return output
 
 
+class MultiConcatModel3(Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.op1 = nn.Conv2d(3, 8, 1)
+        self.op2 = nn.Conv2d(3, 8, 1)
+        self.op3 = nn.Conv2d(3, 8, 1)
+        self.op4 = nn.Conv2d(24, 8, 1)
+        self.op5 = nn.Conv2d(24, 8, 1)
+        self.op6 = nn.Conv2d(24, 8, 1)
+        self.op7 = nn.Conv2d(24, 8, 1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1 = self.op1(x)
+        x2 = self.op2(x)
+        x3 = self.op3(x)
+        cat1 = torch.cat([x1, x2, x3], dim=1)
+        x4 = self.op4(cat1)
+        x5 = self.op5(cat1)
+        x6 = self.op6(cat1)
+        x7 = self.op7(cat1)
+        return torch.cat([x4, x5, x6, x7], dim=1)
+
+
 class ResBlock(Module):
 
     def __init__(self) -> None:
@@ -76,8 +101,11 @@ def forward(self, x: Tensor) -> Tensor:
 
 class ToyCNNPseudoLoss:
 
+    def __init__(self, input_shape=(2, 3, 16, 16)):
+        self.input_shape = input_shape
+
     def __call__(self, model):
-        pseudo_img = torch.rand(2, 3, 16, 16)
+        pseudo_img = torch.rand(self.input_shape)
         pseudo_output = model(pseudo_img)
         return pseudo_output.sum()
 
@@ -157,6 +185,18 @@ def test_trace_multi_cat(self) -> None:
             'op1', 'op2', 'op3'
         ]
 
+        model = MultiConcatModel3()
+        tracer = BackwardTracer(loss_calculator=loss_calculator)
+        path_list = tracer.trace(model)
+        assert len(path_list) == 1
+
+        nonpass2parents = path_list.find_nodes_parents(NONPASS_NODES)
+        assert nonpass2parents['op1'] == list()
+        assert nonpass2parents['op2'] == list()
+        assert nonpass2parents['op3'] == list()
+        assert nonpass2parents['op4'] == nonpass2parents['op5'] == \
+               nonpass2parents['op6'] == nonpass2parents['op7']
+
     def test_repr(self):
         toy_node = PathConvNode('op1')
         assert repr(toy_node) == 'PathConvNode(\'op1\')'

tests/test_core/test_tracer/test_loss_calculator.py

@@ -0,0 +1,25 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from unittest import TestCase
+
+import torch
+from mmengine.hub import get_model
+
+from mmrazor.structures.tracer import (ImageClassifierPseudoLoss,
+                                       SingleStageDetectorPseudoLoss)
+
+
+class TestLossCalculator(TestCase):
+
+    def test_image_classifier_pseudo_loss(self):
+        model = get_model(
+            'mmcls::resnet/resnet34_8xb32_in1k.py', pretrained=False)
+        loss_calculator = ImageClassifierPseudoLoss()
+        loss = loss_calculator(model)
+        assert isinstance(loss, torch.Tensor) and loss.dim() == 0
+
+    def test_single_stage_detector_pseudo_loss(self):
+        model = get_model(
+            'mmdet::retinanet/retinanet_r50_fpn_1x_coco.py', pretrained=False)
+        loss_calculator = SingleStageDetectorPseudoLoss()
+        loss = loss_calculator(model)
+        assert isinstance(loss, torch.Tensor) and loss.dim() == 0