[Refactor]: Unified parameter initialization (#4750)

* Support RetinaNet

* Update RetinaNet init_cfg

* Update RetinaNet init_cfg

* Update model

* Update all model

* Fix type error

* Support specify init_cfg

* Support override init_cfg

* Support ModuleList and Seq

* Use ModuleList

* Update init_cfg

* Add docstr and support Caffe2Xavier

* Update init_weight

* Fix Sequential

* Fix regnet

* Fix BN init_cfg

* Fix error

* Fix unittest

* Fix init error

* Fix bn name

* Fix resnet unittest

* Fix unittest

* Fix ssd and yolact

* Fix layer error

* Fix point_rend

* Fix htc

* Rename init_weight to init_weights

* delete mmcv link

* Fix assert error

* Fix ssd init

* Fix carafe init

* fix lint

* update mmcv version

Co-authored-by: window_hha <123@qq.com>

docs/get_started.md

@@ -11,7 +11,7 @@ The compatible MMDetection and MMCV versions are as below. Please install the co
 
 | MMDetection version |    MMCV version     |
 |:-------------------:|:-------------------:|
-| master              | mmcv-full>=1.2.4, <1.4.0 |
+| master              | mmcv-full>=1.3.2, <1.4.0 |
 | 2.11.0              | mmcv-full>=1.2.4, <1.4.0 |
 | 2.10.0              | mmcv-full>=1.2.4, <1.4.0 |
 | 2.9.0               | mmcv-full>=1.2.4, <1.4.0 |

docs/tutorials/customize_models.md

@@ -32,9 +32,6 @@ class MobileNet(nn.Module):
 
     def forward(self, x):  # should return a tuple
         pass
-
-    def init_weights(self, pretrained=None):
-        pass
 ```
 
 #### 2. Import the module
@@ -157,8 +154,6 @@ class DoubleConvFCBBoxHead(BBoxHead):
         kwargs.setdefault('with_avg_pool', True)
         super(DoubleConvFCBBoxHead, self).__init__(**kwargs)
 
-    def init_weights(self):
-        # conv layers are already initialized by ConvModule
 
     def forward(self, x_cls, x_reg):
 
@@ -186,7 +181,6 @@ class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
 
     def init_mask_head(self, mask_roi_extractor, mask_head):
 
-    def init_weights(self, pretrained):
 
     def forward_dummy(self, x, proposals):
 

mmdet/models/backbones/darknet.py

@@ -1,16 +1,16 @@
 # Copyright (c) 2019 Western Digital Corporation or its affiliates.
 
-import logging
+import warnings
 
 import torch.nn as nn
-from mmcv.cnn import ConvModule, constant_init, kaiming_init
-from mmcv.runner import load_checkpoint
+from mmcv.cnn import ConvModule
+from mmcv.runner import BaseModule
 from torch.nn.modules.batchnorm import _BatchNorm
 
 from ..builder import BACKBONES
 
 
-class ResBlock(nn.Module):
+class ResBlock(BaseModule):
     """The basic residual block used in Darknet. Each ResBlock consists of two
     ConvModules and the input is added to the final output. Each ConvModule is
     composed of Conv, BN, and LeakyReLU. In YoloV3 paper, the first convLayer
@@ -25,14 +25,17 @@ class ResBlock(nn.Module):
             Default: dict(type='BN', requires_grad=True)
         act_cfg (dict): Config dict for activation layer.
             Default: dict(type='LeakyReLU', negative_slope=0.1).
+        init_cfg (dict or list[dict], optional): Initialization config dict.
+            Default: None
     """
 
     def __init__(self,
                  in_channels,
                  conv_cfg=None,
                  norm_cfg=dict(type='BN', requires_grad=True),
-                 act_cfg=dict(type='LeakyReLU', negative_slope=0.1)):
-        super(ResBlock, self).__init__()
+                 act_cfg=dict(type='LeakyReLU', negative_slope=0.1),
+                 init_cfg=None):
+        super(ResBlock, self).__init__(init_cfg)
         assert in_channels % 2 == 0  # ensure the in_channels is even
         half_in_channels = in_channels // 2
 
@@ -53,7 +56,7 @@ def forward(self, x):
 
 
 @BACKBONES.register_module()
-class Darknet(nn.Module):
+class Darknet(BaseModule):
     """Darknet backbone.
 
     Args:
@@ -69,6 +72,9 @@ class Darknet(nn.Module):
         norm_eval (bool): Whether to set norm layers to eval mode, namely,
             freeze running stats (mean and var). Note: Effect on Batch Norm
             and its variants only.
+        pretrained (str, optional): model pretrained path. Default: None
+        init_cfg (dict or list[dict], optional): Initialization config dict.
+            Default: None
 
     Example:
         >>> from mmdet.models import Darknet
@@ -98,10 +104,13 @@ def __init__(self,
                  conv_cfg=None,
                  norm_cfg=dict(type='BN', requires_grad=True),
                  act_cfg=dict(type='LeakyReLU', negative_slope=0.1),
-                 norm_eval=True):
-        super(Darknet, self).__init__()
+                 norm_eval=True,
+                 pretrained=None,
+                 init_cfg=None):
+        super(Darknet, self).__init__(init_cfg)
         if depth not in self.arch_settings:
             raise KeyError(f'invalid depth {depth} for darknet')
+
         self.depth = depth
         self.out_indices = out_indices
         self.frozen_stages = frozen_stages
@@ -122,6 +131,24 @@ def __init__(self,
 
         self.norm_eval = norm_eval
 
+        assert not (init_cfg and pretrained), \
+            'init_cfg and pretrained cannot be setting at the same time'
+        if isinstance(pretrained, str):
+            warnings.warn('DeprecationWarning: pretrained is a deprecated, '
+                          'please use "init_cfg" instead')
+            self.init_cfg = dict(type='Pretrained', checkpoint=pretrained)
+        elif pretrained is None:
+            if init_cfg is None:
+                self.init_cfg = [
+                    dict(type='Kaiming', layer='Conv2d'),
+                    dict(
+                        type='Constant',
+                        val=1,
+                        layer=['_BatchNorm', 'GroupNorm'])
+                ]
+        else:
+            raise TypeError('pretrained must be a str or None')
+
     def forward(self, x):
         outs = []
         for i, layer_name in enumerate(self.cr_blocks):
@@ -132,20 +159,6 @@ def forward(self, x):
 
         return tuple(outs)
 
-    def init_weights(self, pretrained=None):
-        if isinstance(pretrained, str):
-            logger = logging.getLogger()
-            load_checkpoint(self, pretrained, strict=False, logger=logger)
-        elif pretrained is None:
-            for m in self.modules():
-                if isinstance(m, nn.Conv2d):
-                    kaiming_init(m)
-                elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
-                    constant_init(m, 1)
-
-        else:
-            raise TypeError('pretrained must be a str or None')
-
     def _freeze_stages(self):
         if self.frozen_stages >= 0:
             for i in range(self.frozen_stages):

mmdet/models/backbones/detectors_resnet.py

@@ -1,8 +1,13 @@
 import torch.nn as nn
 import torch.utils.checkpoint as cp
-from mmcv.cnn import build_conv_layer, build_norm_layer, constant_init
+from mmcv.cnn import (build_conv_layer, build_norm_layer, constant_init,
+                      kaiming_init)
+from mmcv.runner import Sequential, load_checkpoint
+from torch.nn.modules.batchnorm import _BatchNorm
 
+from mmdet.utils import get_root_logger
 from ..builder import BACKBONES
+from .resnet import BasicBlock
 from .resnet import Bottleneck as _Bottleneck
 from .resnet import ResNet
 
@@ -22,6 +27,8 @@ class Bottleneck(_Bottleneck):
              added for ``rfp_feat``. Otherwise, the structure is the same as
              base class.
          sac (dict, optional): Dictionary to construct SAC. Default: None.
+         init_cfg (dict or list[dict], optional): Initialization config dict.
+            Default: None
     """
     expansion = 4
 
@@ -30,8 +37,10 @@ def __init__(self,
                  planes,
                  rfp_inplanes=None,
                  sac=None,
+                 init_cfg=None,
                  **kwargs):
-        super(Bottleneck, self).__init__(inplanes, planes, **kwargs)
+        super(Bottleneck, self).__init__(
+            inplanes, planes, init_cfg=init_cfg, **kwargs)
 
         assert sac is None or isinstance(sac, dict)
         self.sac = sac
@@ -56,12 +65,9 @@ def __init__(self,
                 1,
                 stride=1,
                 bias=True)
-        self.init_weights()
-
-    def init_weights(self):
-        """Initialize the weights."""
-        if self.rfp_inplanes:
-            constant_init(self.rfp_conv, 0)
+            if init_cfg is None:
+                self.init_cfg = dict(
+                    type='Constant', val=0, override=dict(name='rfp_conv'))
 
     def rfp_forward(self, x, rfp_feat):
         """The forward function that also takes the RFP features as input."""
@@ -110,7 +116,7 @@ def _inner_forward(x):
         return out
 
 
-class ResLayer(nn.Sequential):
+class ResLayer(Sequential):
     """ResLayer to build ResNet style backbone for RPF in detectoRS.
 
     The difference between this module and base class is that we pass
@@ -216,7 +222,6 @@ class DetectoRS_ResNet(ResNet):
             base class.
         output_img (bool): If ``True``, the input image will be inserted into
             the starting position of output. Default: False.
-        pretrained (str, optional): The pretrained model to load.
     """
 
     arch_settings = {
@@ -231,12 +236,15 @@ def __init__(self,
                  rfp_inplanes=None,
                  output_img=False,
                  pretrained=None,
+                 init_cfg=None,
                  **kwargs):
+        assert init_cfg is None, 'To prevent abnormal initialization ' \
+                                 'behavior, init_cfg is not allowed to be set'
+        self.pretrained = pretrained
         self.sac = sac
         self.stage_with_sac = stage_with_sac
         self.rfp_inplanes = rfp_inplanes
         self.output_img = output_img
-        self.pretrained = pretrained
         super(DetectoRS_ResNet, self).__init__(**kwargs)
 
         self.inplanes = self.stem_channels
@@ -274,6 +282,36 @@ def __init__(self,
 
         self._freeze_stages()
 
+    # In order to be properly initialized by RFP
+    def init_weights(self):
+        # Calling this method will cause parameter initialization exception
+        # super(DetectoRS_ResNet, self).init_weights()
+
+        if isinstance(self.pretrained, str):
+            logger = get_root_logger()
+            load_checkpoint(self, self.pretrained, strict=False, logger=logger)
+        elif self.pretrained is None:
+            for m in self.modules():
+                if isinstance(m, nn.Conv2d):
+                    kaiming_init(m)
+                elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
+                    constant_init(m, 1)
+
+            if self.dcn is not None:
+                for m in self.modules():
+                    if isinstance(m, Bottleneck) and hasattr(
+                            m.conv2, 'conv_offset'):
+                        constant_init(m.conv2.conv_offset, 0)
+
+            if self.zero_init_residual:
+                for m in self.modules():
+                    if isinstance(m, Bottleneck):
+                        constant_init(m.norm3, 0)
+                    elif isinstance(m, BasicBlock):
+                        constant_init(m.norm2, 0)
+        else:
+            raise TypeError('pretrained must be a str or None')
+
     def make_res_layer(self, **kwargs):
         """Pack all blocks in a stage into a ``ResLayer`` for DetectoRS."""
         return ResLayer(**kwargs)

mmdet/models/backbones/hourglass.py

@@ -1,12 +1,13 @@
 import torch.nn as nn
 from mmcv.cnn import ConvModule
+from mmcv.runner import BaseModule
 
 from ..builder import BACKBONES
 from ..utils import ResLayer
 from .resnet import BasicBlock
 
 
-class HourglassModule(nn.Module):
+class HourglassModule(BaseModule):
     """Hourglass Module for HourglassNet backbone.
 
     Generate module recursively and use BasicBlock as the base unit.
@@ -18,14 +19,17 @@ class HourglassModule(nn.Module):
         stage_blocks (list[int]): Number of sub-modules stacked in current and
             follow-up HourglassModule.
         norm_cfg (dict): Dictionary to construct and config norm layer.
+        init_cfg (dict or list[dict], optional): Initialization config dict.
+            Default: None
     """
 
     def __init__(self,
                  depth,
                  stage_channels,
                  stage_blocks,
-                 norm_cfg=dict(type='BN', requires_grad=True)):
-        super(HourglassModule, self).__init__()
+                 norm_cfg=dict(type='BN', requires_grad=True),
+                 init_cfg=None):
+        super(HourglassModule, self).__init__(init_cfg)
 
         self.depth = depth
 
@@ -78,7 +82,7 @@ def forward(self, x):
 
 
 @BACKBONES.register_module()
-class HourglassNet(nn.Module):
+class HourglassNet(BaseModule):
     """HourglassNet backbone.
 
     Stacked Hourglass Networks for Human Pose Estimation.
@@ -95,6 +99,9 @@ class HourglassNet(nn.Module):
             HourglassModule.
         feat_channel (int): Feature channel of conv after a HourglassModule.
         norm_cfg (dict): Dictionary to construct and config norm layer.
+        pretrained (str, optional): model pretrained path. Default: None
+        init_cfg (dict or list[dict], optional): Initialization config dict.
+            Default: None
 
     Example:
         >>> from mmdet.models import HourglassNet
@@ -115,8 +122,12 @@ def __init__(self,
                  stage_channels=(256, 256, 384, 384, 384, 512),
                  stage_blocks=(2, 2, 2, 2, 2, 4),
                  feat_channel=256,
-                 norm_cfg=dict(type='BN', requires_grad=True)):
-        super(HourglassNet, self).__init__()
+                 norm_cfg=dict(type='BN', requires_grad=True),
+                 pretrained=None,
+                 init_cfg=None):
+        assert init_cfg is None, 'To prevent abnormal initialization ' \
+                                 'behavior, init_cfg is not allowed to be set'
+        super(HourglassNet, self).__init__(init_cfg)
 
         self.num_stacks = num_stacks
         assert self.num_stacks >= 1
@@ -161,17 +172,10 @@ def __init__(self,
 
         self.relu = nn.ReLU(inplace=True)
 
-    def init_weights(self, pretrained=None):
-        """Init module weights.
-
-        We do nothing in this function because all modules we used
-        (ConvModule, BasicBlock and etc.) have default initialization, and
-        currently we don't provide pretrained model of HourglassNet.
-
-        Detector's __init__() will call backbone's init_weights() with
-        pretrained as input, so we keep this function.
-        """
+    def init_weights(self):
+        """Init module weights."""
         # Training Centripetal Model needs to reset parameters for Conv2d
+        super(HourglassNet, self).init_weights()
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
                 m.reset_parameters()