[Refactor] Move res-layer to models.utils (#537)

* move reslayer to utils

* update configs

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)

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)

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)

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'
 ]

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