From 560446af852b7f0b268e9293e38ddc182d5f1e16 Mon Sep 17 00:00:00 2001
From: Yixiao Fang <36138628+fangyixiao18@users.noreply.github.com>
Date: Fri, 28 Oct 2022 10:10:15 +0800
Subject: [PATCH] [Refactor] Move res-layer to models.utils (#537)

* move reslayer to utils

* update configs
---
 .../coco/mask_rcnn_r50_c4_mstrain_1x_coco.py  |  7 ++---
 .../coco/mask_rcnn_r50_c4_mstrain_2x_coco.py  |  7 ++---
 .../faster_rcnn_r50_c4_mstrain_24k_voc0712.py |  7 ++---
 mmselfsup/models/utils/__init__.py            |  8 ++++-
 .../models/utils/res_layer_extra_norm.py      | 31 +++++++++++++++++++
 .../mmdetection/res_layer_extra_norm.py       | 27 ----------------
 6 files changed, 47 insertions(+), 40 deletions(-)
 create mode 100644 mmselfsup/models/utils/res_layer_extra_norm.py
 delete mode 100644 tools/benchmarks/mmdetection/res_layer_extra_norm.py

diff --git a/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_1x_coco.py b/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_1x_coco.py
index 0e62600c8..cd531f8d4 100644
--- a/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_1x_coco.py
+++ b/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_1x_coco.py
@@ -4,6 +4,9 @@
     '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py'
 ]
 
+custom_imports = dict(
+    imports=['mmselfsup.models.utils.res_layer_extra_norm'],
+    allow_failed_imports=False)
 norm_cfg = dict(type='SyncBN', requires_grad=True)
 model = dict(
     backbone=dict(frozen_stages=-1, norm_cfg=norm_cfg, norm_eval=False),
@@ -30,7 +33,3 @@
 ]
 
 data = dict(train=dict(pipeline=train_pipeline))
-
-custom_imports = dict(
-    imports=['tools.benchmarks.mmdetection.res_layer_extra_norm'],
-    allow_failed_imports=False)
diff --git a/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_2x_coco.py b/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_2x_coco.py
index 82c8d3020..f0b7c6bec 100644
--- a/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_2x_coco.py
+++ b/configs/benchmarks/mmdetection/coco/mask_rcnn_r50_c4_mstrain_2x_coco.py
@@ -4,6 +4,9 @@
     '../_base_/schedules/schedule_2x.py', '../_base_/default_runtime.py'
 ]
 
+custom_imports = dict(
+    imports=['mmselfsup.models.utils.res_layer_extra_norm'],
+    allow_failed_imports=False)
 norm_cfg = dict(type='SyncBN', requires_grad=True)
 model = dict(
     backbone=dict(frozen_stages=-1, norm_cfg=norm_cfg, norm_eval=False),
@@ -30,7 +33,3 @@
 ]
 
 data = dict(train=dict(pipeline=train_pipeline))
-
-custom_imports = dict(
-    imports=['tools.benchmarks.mmdetection.res_layer_extra_norm'],
-    allow_failed_imports=False)
diff --git a/configs/benchmarks/mmdetection/voc0712/faster_rcnn_r50_c4_mstrain_24k_voc0712.py b/configs/benchmarks/mmdetection/voc0712/faster_rcnn_r50_c4_mstrain_24k_voc0712.py
index 969c27b06..f8d978f85 100644
--- a/configs/benchmarks/mmdetection/voc0712/faster_rcnn_r50_c4_mstrain_24k_voc0712.py
+++ b/configs/benchmarks/mmdetection/voc0712/faster_rcnn_r50_c4_mstrain_24k_voc0712.py
@@ -3,6 +3,9 @@
     '../_base_/schedules/schedule_24k.py', '../_base_/default_runtime.py'
 ]
 
+custom_imports = dict(
+    imports=['mmselfsup.models.utils.res_layer_extra_norm'],
+    allow_failed_imports=False)
 norm_cfg = dict(type='SyncBN', requires_grad=True)
 model = dict(
     backbone=dict(frozen_stages=-1, norm_cfg=norm_cfg, norm_eval=False),
@@ -78,7 +81,3 @@
         dict(type='TextLoggerHook', by_epoch=False),
         # dict(type='TensorboardLoggerHook')
     ])
-
-custom_imports = dict(
-    imports=['tools.benchmarks.mmdetection.res_layer_extra_norm'],
-    allow_failed_imports=False)
diff --git a/mmselfsup/models/utils/__init__.py b/mmselfsup/models/utils/__init__.py
index 55fb6982a..b3fad9a35 100644
--- a/mmselfsup/models/utils/__init__.py
+++ b/mmselfsup/models/utils/__init__.py
@@ -11,9 +11,15 @@
 from .transformer_blocks import (CAETransformerRegressorLayer,
                                  MultiheadAttention, TransformerEncoderLayer)
 
+try:
+    from .res_layer_extra_norm import ResLayerExtraNorm
+except ImportError:
+    ResLayerExtraNorm = None
+
 __all__ = [
     'Accuracy', 'accuracy', 'ExtractProcess', 'MultiExtractProcess',
     'GatherLayer', 'knn_classifier', 'MultiPooling', 'MultiPrototypes',
     'build_2d_sincos_position_embedding', 'Sobel', 'MultiheadAttention',
-    'TransformerEncoderLayer', 'CAETransformerRegressorLayer', 'Encoder'
+    'TransformerEncoderLayer', 'CAETransformerRegressorLayer', 'Encoder',
+    'ResLayerExtraNorm'
 ]
diff --git a/mmselfsup/models/utils/res_layer_extra_norm.py b/mmselfsup/models/utils/res_layer_extra_norm.py
new file mode 100644
index 000000000..9d9105a60
--- /dev/null
+++ b/mmselfsup/models/utils/res_layer_extra_norm.py
@@ -0,0 +1,31 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from mmcv.cnn import build_norm_layer
+from mmcv.runner import auto_fp16
+
+try:
+    from mmdet.models.backbones import ResNet
+    from mmdet.models.builder import SHARED_HEADS
+    from mmdet.models.roi_heads.shared_heads.res_layer import ResLayer
+
+    @SHARED_HEADS.register_module()
+    class ResLayerExtraNorm(ResLayer):
+
+        def __init__(self, *args, **kwargs):
+            super(ResLayerExtraNorm, self).__init__(*args, **kwargs)
+
+            block = ResNet.arch_settings[kwargs['depth']][0]
+            self.add_module(
+                'norm',
+                build_norm_layer(self.norm_cfg,
+                                 64 * 2**self.stage * block.expansion)[1])
+
+        @auto_fp16()
+        def forward(self, x):
+            res_layer = getattr(self, f'layer{self.stage + 1}')
+            norm = getattr(self, 'norm')
+            x = res_layer(x)
+            out = norm(x)
+            return out
+
+except ImportError:
+    pass
diff --git a/tools/benchmarks/mmdetection/res_layer_extra_norm.py b/tools/benchmarks/mmdetection/res_layer_extra_norm.py
deleted file mode 100644
index 0364ba427..000000000
--- a/tools/benchmarks/mmdetection/res_layer_extra_norm.py
+++ /dev/null
@@ -1,27 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-from mmcv.cnn import build_norm_layer
-from mmcv.runner import auto_fp16
-from mmdet.models.backbones import ResNet
-from mmdet.models.builder import SHARED_HEADS
-from mmdet.models.roi_heads.shared_heads.res_layer import ResLayer
-
-
-@SHARED_HEADS.register_module()
-class ResLayerExtraNorm(ResLayer):
-
-    def __init__(self, *args, **kwargs):
-        super(ResLayerExtraNorm, self).__init__(*args, **kwargs)
-
-        block = ResNet.arch_settings[kwargs['depth']][0]
-        self.add_module(
-            'norm',
-            build_norm_layer(self.norm_cfg,
-                             64 * 2**self.stage * block.expansion)[1])
-
-    @auto_fp16()
-    def forward(self, x):
-        res_layer = getattr(self, f'layer{self.stage + 1}')
-        norm = getattr(self, 'norm')
-        x = res_layer(x)
-        out = norm(x)
-        return out