[Features] Add logger for initialization of parameters

mmcv/cnn/utils/weight_init.py

@@ -8,6 +8,7 @@
 import torch.nn as nn
 from torch import Tensor
 
+from mmcv.runner.base_module import update_init_infos
 from mmcv.utils import Registry, build_from_cfg, get_logger, print_log
 
 INITIALIZERS = Registry('initializer')
@@ -122,6 +123,10 @@ def __init__(self, *, bias=0, bias_prob=None, layer=None):
             self.bias = bias
         self.layer = [layer] if isinstance(layer, str) else layer
 
+    def _get_init_info(self):
+        info = f', bias={self.bias}'
+        return info
+
 
 @INITIALIZERS.register_module(name='Constant')
 class ConstantInit(BaseInit):
@@ -152,6 +157,12 @@ def init(m):
                     constant_init(m, self.val, self.bias)
 
         module.apply(init)
+        if hasattr(module, 'params_init_info'):
+            update_init_infos(module, init_info=self._get_init_info())
+
+    def _get_init_info(self):
+        info = f'{self.__class__.__name__}: val={self.val}'
+        return info + super()._get_init_info()
 
 
 @INITIALIZERS.register_module(name='Xavier')
@@ -189,6 +200,13 @@ def init(m):
                     xavier_init(m, self.gain, self.bias, self.distribution)
 
         module.apply(init)
+        if hasattr(module, 'params_init_info'):
+            update_init_infos(module, init_info=self._get_init_info())
+
+    def _get_init_info(self):
+        info = f'{self.__class__.__name__}: gain={self.gain}, ' \
+               f'distribution={self.distribution}'
+        return info + super()._get_init_info()
 
 
 @INITIALIZERS.register_module(name='Normal')
@@ -225,6 +243,12 @@ def init(m):
                     normal_init(m, self.mean, self.std, self.bias)
 
         module.apply(init)
+        if hasattr(module, 'params_init_info'):
+            update_init_infos(module, init_info=self._get_init_info())
+
+    def _get_init_info(self):
+        info = f'{self.__class__.__name__}: mean={self.mean}, std={self.std}'
+        return info + super()._get_init_info()
 
 
 @INITIALIZERS.register_module(name='TruncNormal')
@@ -273,6 +297,13 @@ def init(m):
                                       self.bias)
 
         module.apply(init)
+        if hasattr(module, 'params_init_info'):
+            update_init_infos(module, init_info=self._get_init_info())
+
+    def _get_init_info(self):
+        info = f'{self.__class__.__name__}: a={self.a}, b={self.b},' \
+               f' mean={self.mean}, std={self.std}'
+        return info + super()._get_init_info()
 
 
 @INITIALIZERS.register_module(name='Uniform')
@@ -309,6 +340,12 @@ def init(m):
                     uniform_init(m, self.a, self.b, self.bias)
 
         module.apply(init)
+        if hasattr(module, 'params_init_info'):
+            update_init_infos(module, init_info=self._get_init_info())
+
+    def _get_init_info(self):
+        info = f'{self.__class__.__name__}: a={self.a}, b={self.b}'
+        return info + super()._get_init_info()
 
 
 @INITIALIZERS.register_module(name='Kaiming')
@@ -364,6 +401,14 @@ def init(m):
                                  self.bias, self.distribution)
 
         module.apply(init)
+        if hasattr(module, 'params_init_info'):
+            update_init_infos(module, init_info=self._get_init_info())
+
+    def _get_init_info(self):
+        info = f'{self.__class__.__name__}: a={self.a}, mode={self.mode}, ' \
+               f'nonlinearity={self.nonlinearity}, ' \
+               f'distribution ={self.distribution}'
+        return info + super()._get_init_info()
 
 
 @INITIALIZERS.register_module(name='Caffe2Xavier')
@@ -422,6 +467,13 @@ def __call__(self, module):
                 self.prefix, self.checkpoint, map_location=self.map_location)
             load_state_dict(module, state_dict, strict=False, logger=logger)
 
+        if hasattr(module, 'params_init_info'):
+            update_init_infos(module, init_info=self._get_init_info())
+
+    def _get_init_info(self):
+        info = f'{self.__class__.__name__}: load from {self.checkpoint}'
+        return info
+
 
 def _initialize(module, cfg, wholemodule=False):
     func = build_from_cfg(cfg, INITIALIZERS)

mmcv/runner/base_module.py

@@ -1,10 +1,18 @@
 # Copyright (c) Open-MMLab. All rights reserved.
 import warnings
 from abc import ABCMeta
+from collections import defaultdict
 
 import torch.nn as nn
 
-from mmcv import ConfigDict
+from mmcv.utils.logging import logger_initialized, print_log
+
+
+def update_init_infos(module, *, init_info):
+    for param in module.parameters():
+        if module.params_init_info[param]['tmp_sum_value'] != param.data.sum():
+            module.params_init_info[param]['init_info'] = init_info
+            module.params_init_info[param]['tmp_sum_value'] = param.data.sum()
 
 
 class BaseModule(nn.Module, metaclass=ABCMeta):
@@ -26,6 +34,14 @@ def __init__(self, init_cfg=None):
         self._is_init = False
         self.init_cfg = init_cfg
 
+        # The `params_init_info` is used to record the initialization
+        # information of the parameters
+        # the key should be the obj:`nn.Parameter` of model and the value
+        # should be a dict contains
+        # `params_name`, `init_info` and `tmp_sum_value`.
+        # this attribute would be deleted after all parameters is initialized.
+        self.params_init_info = defaultdict(dict)
+
         # Backward compatibility in derived classes
         # if pretrained is not None:
         #     warnings.warn('DeprecationWarning: pretrained is a deprecated \
@@ -38,12 +54,33 @@ def is_init(self):
 
     def init_weights(self):
         """Initialize the weights."""
-        from ..cnn import initialize
 
+        # judge if it is Topmost module
+        top_most_module = len(self.params_init_info) == 0
+        if top_most_module:
+            for name, param in self.named_parameters():
+                self.params_init_info[param]['params_name'] = name
+                self.params_init_info[param][
+                    'init_info'] = 'The value is the same as ' \
+                                   'the default initialization of PyTorch'
+                self.params_init_info[param]['tmp_sum_value'] = param.data.sum(
+                )
+            # pass `params_init_info` to all submodules
+            for sub_moduls in self.modules():
+                sub_moduls.params_init_info = self.params_init_info
+
+        loggernames = list(logger_initialized.keys())
+        loggername = loggernames[0] if len(loggernames) > 0 else 'mmcv'
+
+        from ..cnn import initialize
+        modulename = self.__class__.__name__
         if not self._is_init:
             if self.init_cfg:
+                print_log(
+                    f'initialize {modulename} with init_cfg {self.init_cfg}',
+                    logger=loggername)
                 initialize(self, self.init_cfg)
-                if isinstance(self.init_cfg, (dict, ConfigDict)):
+                if isinstance(self.init_cfg, dict):
                     # Avoid the parameters of the pre-training model
                     # being overwritten by the init_weights
                     # of the children.
@@ -53,11 +90,24 @@ def init_weights(self):
             for m in self.children():
                 if hasattr(m, 'init_weights'):
                     m.init_weights()
+                    # user may overload the `init_weights`
+                    update_init_infos(
+                        m,
+                        init_info='Initialized by user-defined `init_weights`')
+
             self._is_init = True
         else:
             warnings.warn(f'init_weights of {self.__class__.__name__} has '
                           f'been called more than once.')
 
+        if top_most_module:
+            for item in list(self.params_init_info.values()):
+                print_log(
+                    f"{item['params_name']} - {item['init_info']}",
+                    logger=loggername)
+            for sub_moduls in self.modules():
+                del sub_moduls.params_init_info
+
     def __repr__(self):
         s = super().__repr__()
         if self.init_cfg: