mmdetection代码阅读系列（六）：BaseModule的初始化

BaseModule初始化的参数init_cfg例子如下

dict(
    type='Normal',
    layer='Conv2d',
    std=0.01,
    override=dict(
        type='Normal',
        name='conv_cls',
        std=0.01,
        bias_prob=0.01)
)

type: (str) 用于初始化的INITIALIZERS的name
layer:(list[str], str) 需要对当前module及其children哪些类型的层进行初始化
bias_prob/bias:(float) 初始化的bias的值
override: 对指定属性模块进行重新的overrid.
name: (str) 就是对当前模块的哪个属性模块使用新的初始化方法进行override
std: NormalInit INITIALIZERS 的参数

override

init_cfg只需要定义在最外面的model上就可以，如果内部某些子模块初始化方法不同于model，则可以在子模块中定义init_cfg或者指定model的init_cfg里的override属性

比如如下代码，B的参数中定义的init_cfg会作用于m1，m2, m3，因为INITIALIZERS 会调用b.apply(init_func)，apply会将自己和children都使用init_func处理。

class A(BaseModule):
	def __init__(self, init_cfg=None):
		super().__init__(init_cfg=init_cfg)
		self.c1 = nn.Conv2d(...)
		self.c2 = nn.Conv2d(...)

class B(BaseModule):
	def __init__(self, init_cfg):
		super().__init__(init_cfg)
		self.m1 = A()
		self.m2 = A()
		self.m3 = A()

如果，m1和m2,m3的初始方法不一样，可以使用上面override,也可以在定义的时候给定

B(
dict(
    type='Normal',
    layer='Conv2d',
    std=0.01,
    override=dict(
        type='Normal',
        name='m1',
        std=0.01,
        bias_prob=0.01)
)
或者
class B(BaseModule):
	def __init__(self, init_cfg):
		super().__init__(init_cfg)
		self.m1 = A(dict(
        type='Conv2d',
        name='m1',
        std=0.01,
        bias_prob=0.01))
		self.m2 = A()
		self.m3 = A()

BaseModule 初始化调用关系

调用initialize ->_initialize_override/_initialize --> INITIALIZERS进行初始化

class BaseModule(nn.Moudle):
    def init_weights(self):
    	...
    	if not self._is_init and self.init_cfg::
	        initialize(self, self.init_cfg)  # 按照初始化自己以及children
	        for m in self.children():
	        	if hasattr(m, 'init_weights'): # 如果children自己定义了init_cfg，则使用自己的初始化
	        		m.init_weights()
        ...


# mmcv\cnn\utils\weight_init.py

def initialize(module, init_cfg):
    if isinstance(init_cfg, dict):
        init_cfg = [init_cfg]
    for cfg in init_cfg:
        cp_cfg = copy.deepcopy(cfg)
        override = cp_cfg.pop('override', None)
        _initialize(module, cp_cfg)

        if override is not None:
            cp_cfg.pop('layer', None)
            _initialize_override(module, override, cp_cfg)

def _initialize(module, cfg, wholemodule=False):
    func = build_from_cfg(cfg, INITIALIZERS)
    func.wholemodule = wholemodule  # for _initialize_override
    func(module)

def _initialize_override(module, override, cfg):
    override = [override] if isinstance(override, dict) else override
    for override_ in override:
        cp_override = copy.deepcopy(override_)
        name = cp_override.pop('name', None)
        # if override only has name kay, it means use args in init_cfg
        if not cp_override:
            cp_override.update(cfg)
            
        if hasattr(module, name):
            _initialize(getattr(module, name), cp_override, wholemodule=True)
    ...

INITIALIZERS

• bias / bias_prob: float, 表示的是将bias初始化那个值，两个设置一个就行，bias_prob的优先级更高。
• layer: (str, list)，表示需要对哪些类型的层进行初始化，比如’Conv2d’，在__call__里面使用：

  layername = m.__class__.__name__
  basesname = _get_bases_name(m)
  if len(set(self.layer) & set([layername] + basesname)):
      normal_init(m, self.mean, self.std, self.bias)
  

  set([layername] + basesname) 是获得模块m的类名以及所有的祖先类名，和set(self.layer)求交集(&)，如果不为空，表明m是需要初始化的类型。
• __call__函数输入的module是nn.Module，module.apply(func)的实现发现它会现将self.children进行func处理，然后，再func(self)
• 综合上面可知，INITIALIZERS会将module以及module的children中符合layer类型的模块进行初始化。

# mmcv\cnn\utils\weight_init.py
class BaseInit(object):
    def __init__(self, *, bias=0, bias_prob=None, layer=None):
        self.wholemodule = False
        ...
        layer = []
        if bias_prob is not None:
            self.bias = bias_init_with_prob(bias_prob)
        else:
            self.bias = bias
        self.layer = [layer] if isinstance(layer, str) else layer

@INITIALIZERS.register_module(name='Normal')
class NormalInit(BaseInit):
    def __call__(self, module):
        def init(m):
            if self.wholemodule:
                normal_init(m, self.mean, self.std, self.bias)
            else:
                layername = m.__class__.__name__
                basesname = _get_bases_name(m)
                if len(set(self.layer) & set([layername] + basesname)):
                    normal_init(m, self.mean, self.std, self.bias)

        module.apply(init)

# torch\nn\modules\module.py
class Module(object):
	...
    def apply(self, fn):
        for module in self.children():
            module.apply(fn)
        fn(self)
        return self
    ...