Repository: open-mmlab/mmengine

Issue:
Unexpected weight initialization
Thanks for your error report and we appreciate it a lot.

**Checklist**

1. I have searched related issues but cannot get the expected help.
2. I have read the [FAQ documentation](https://mmengine.readthedocs.io/en/latest/faq.html) but cannot get the expected help.
3. The bug has not been fixed in the latest version.

**Describe the bug**

When using `Pretrained` init_cfg to init a model, the outside init_cfg will **disable all submodules' init_weights**.

If the outside model's pretrained weight does not strictly match, e.g. different num_classes, the mismatched submodule's weight will initialize **by pytorch's default init weight logic**, instead of the submodule's init_cfg. 

This lead's to some unacceptable results, e.g. a huge classification loss due to the failed initialization of the classification head when finetuning a model.


**Reproduction**

In mmdet, finetuning rtmdet_l_8xb32-300e_coco

```python
_base_ = './rtmdet_l_8xb32-300e_coco.py'

checkpoint = 'work_dir/rtmdet/rtmdet_l_8xb32-300e_coco/epoch_300.pth'
model = dict(bbox_head=dict(num_classes=1), init_cfg=dict(type='Pretrained', checkpoint=checkpoint))
```




Discussion:
How about adding an `overwriten` argument to `PretrainedInit`? If `overwriten=False`(default value), `_is_init` of all submodules will be set to `True` during loading the pre-trained model, otherwise submodules could overwrite the loaded weight.

Resolving patch:
diff --git a/mmengine/model/base_module.py b/mmengine/model/base_module.py
--- a/mmengine/model/base_module.py
+++ b/mmengine/model/base_module.py
@@ -5,7 +5,7 @@
 from abc import ABCMeta
 from collections import defaultdict
 from logging import FileHandler
-from typing import Iterable, Optional
+from typing import Iterable, List, Optional, Union
 
 import torch.nn as nn
 
@@ -26,11 +26,17 @@ class BaseModule(nn.Module, metaclass=ABCMeta):
     - ``_params_init_info``: Used to track the parameter initialization
       information. This attribute only exists during executing the
       ``init_weights``.
+
+    Note:
+        :obj:`PretrainedInit` has a higher priority than any other
+        initializer. The loaded pretrained weights will overwrite
+        the previous initialized weights.
+
     Args:
-        init_cfg (dict, optional): Initialization config dict.
+        init_cfg (dict or List[dict], optional): Initialization config dict.
     """
 
-    def __init__(self, init_cfg=None):
+    def __init__(self, init_cfg: Union[dict, List[dict], None] = None):
         """Initialize BaseModule, inherited from `torch.nn.Module`"""
 
         # NOTE init_cfg can be defined in different levels, but init_cfg
@@ -100,14 +106,25 @@ def init_weights(self):
                     f'initialize {module_name} with init_cfg {self.init_cfg}',
                     logger=logger_name,
                     level=logging.DEBUG)
-                initialize(self, self.init_cfg)
+
+                init_cfgs = self.init_cfg
                 if isinstance(self.init_cfg, dict):
-                    # prevent the parameters of
-                    # the pre-trained model
-                    # from being overwritten by
-                    # the `init_weights`
-                    if self.init_cfg['type'] == 'Pretrained':
-                        return
+                    init_cfgs = [self.init_cfg]
+
+                # PretrainedInit has higher priority than any other init_cfg.
+                # Therefore we initialize `pretrained_cfg` last to overwrite
+                # the previous initialized weights.
+                # See details in https://github.com/open-mmlab/mmengine/issues/691 # noqa E501
+                other_cfgs = []
+                pretrained_cfg = []
+                for init_cfg in init_cfgs:
+                    assert isinstance(init_cfg, dict)
+                    if init_cfg['type'] == 'Pretrained':
+                        pretrained_cfg.append(init_cfg)
+                    else:
+                        other_cfgs.append(init_cfg)
+
+                initialize(self, other_cfgs)
 
             for m in self.children():
                 if hasattr(m, 'init_weights'):
@@ -118,7 +135,8 @@ def init_weights(self):
                         init_info=f'Initialized by '
                         f'user-defined `init_weights`'
                         f' in {m.__class__.__name__} ')
-
+            if self.init_cfg and pretrained_cfg:
+                initialize(self, pretrained_cfg)
             self._is_init = True
         else:
             warnings.warn(f'init_weights of {self.__class__.__name__} has '


Test patch:
diff --git a/tests/test_model/test_base_module.py b/tests/test_model/test_base_module.py
--- a/tests/test_model/test_base_module.py
+++ b/tests/test_model/test_base_module.py
@@ -1,9 +1,13 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import copy
 import logging
+import os.path as osp
+import tempfile
 from unittest import TestCase
 
 import torch
 from torch import nn
+from torch.nn.init import constant_
 
 from mmengine.logging.logger import MMLogger
 from mmengine.model import BaseModule, ModuleDict, ModuleList, Sequential
@@ -90,6 +94,7 @@ def __init__(self,
 class TestBaseModule(TestCase):
 
     def setUp(self) -> None:
+        self.temp_dir = tempfile.TemporaryDirectory()
         self.BaseModule = BaseModule()
         self.model_cfg = dict(
             type='FooModel',
@@ -110,6 +115,7 @@ def setUp(self) -> None:
         self.logger = MMLogger.get_instance(self._testMethodName)
 
     def tearDown(self) -> None:
+        self.temp_dir.cleanup()
         logging.shutdown()
         MMLogger._instance_dict.clear()
         return super().tearDown()
@@ -177,6 +183,45 @@ def test_init_weights(self):
         assert torch.equal(self.model.reg.bias,
                            torch.full(self.model.reg.bias.shape, 2.0))
 
+        # Test build model from Pretrained weights
+
+        class CustomLinear(BaseModule):
+
+            def __init__(self, init_cfg=None):
+                super().__init__(init_cfg)
+                self.linear = nn.Linear(1, 1)
+
+            def init_weights(self):
+                constant_(self.linear.weight, 1)
+                constant_(self.linear.bias, 2)
+
+        @FOOMODELS.register_module()
+        class PratrainedModel(FooModel):
+
+            def __init__(self,
+                         component1=None,
+                         component2=None,
+                         component3=None,
+                         component4=None,
+                         init_cfg=None) -> None:
+                super().__init__(component1, component2, component3,
+                                 component4, init_cfg)
+                self.linear = CustomLinear()
+
+        checkpoint_path = osp.join(self.temp_dir.name, 'test.pth')
+        torch.save(self.model.state_dict(), checkpoint_path)
+        model_cfg = copy.deepcopy(self.model_cfg)
+        model_cfg['type'] = 'PratrainedModel'
+        model_cfg['init_cfg'] = dict(
+            type='Pretrained', checkpoint=checkpoint_path)
+        model = FOOMODELS.build(model_cfg)
+        ori_layer_weight = model.linear.linear.weight.clone()
+        ori_layer_bias = model.linear.linear.bias.clone()
+        model.init_weights()
+
+        self.assertTrue((ori_layer_weight != model.linear.linear.weight).any())
+        self.assertTrue((ori_layer_bias != model.linear.linear.bias).any())
+
     def test_dump_init_info(self):
         import os
         import shutil