MMLab中自定义模块初始化方法

这方面好像介绍的少,看了下基本原理,下面介绍下如何在搭载权重的模型中加入自定义模块时方便的进行初始化。

MMLab的逻辑时对每个部分进行初始化,若此部分定义了初始化方法为Pretrained则加载权重,然后对内部其他模块不再进行其他的初始化操作。

但其存在着一个问题。例如,需要对backbone改进,加入自定义模块后同时需要原模型的预训练权重,此时无法方便的对新加入模块进行初始化操作(因为代码决定其会跳过了此部分初始化,直接进行下一部分的初始化操作)。源代码如下(在BaseModule)中:

        if not self._is_init:
            if self.init_cfg:
                print_log(
                    f'initialize {module_name} with init_cfg {self.init_cfg}',
                    logger=logger_name)
                initialize(self, 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

            for m in self.children():
                if hasattr(m, 'init_weights'):
                    m.init_weights()
                    # users may overload the `init_weights`
                    update_init_info(
                        m,
                        init_info=f'Initialized by '
                        f'user-defined `init_weights`'
                        f' in {m.__class__.__name__} ')

            self._is_init = True

那么如何对自定义模块方便地进行初始化呢,下面介绍三种方法:

(1)定义一个my_weight_init()对自定义模块中的所有module进行初始化操作,其优点是可操作性强,但设置复杂。代码如下:

def my_module_weights_init(target_module):
    for m in target_module.modules():
        if type(m) == nn.Conv2d:
            nn.init.xavier_normal_(m.weight.data)
            nn.init.constant_(m.bias.data, 0.0)

对自定义的模块的初始化直接调用apply即可。

(2)对于MMLab中定义好的模块,若其存在init_cfg则可直接输入相关设置参数进行初始化操作。

(3)最为方便的方法,在mmcv.cnn.utils.weight_init中存在initialize函数,可通过相关参数对函数内部所有相关层进行初始化操作,主要原理是建立初始化器的实例化对象,对模块参数进行处理。mmcv中目前可调用一下八种方法进行初始化,位于mmcv.cnn.utilsz中。

'ConstantInit', 'XavierInit', 'NormalInit', 'TruncNormalInit', 'UniformInit', 'KaimingInit', 'PretrainedInit', 'Caffe2XavierInit'

上述initialize方法的相关代码如下:

def my_module_weights_init(target_module, init_cfg):
    from mmcv.cnn.utils.weight_init import initialize
    initialize(target_module, init_cfg)

调用初始化方法的代码(可直接调用initialize方法,我为了方便好看改了个名):

if self.training:
    my_module_init_cfg = [dict(type='TruncNormal', layer=['Conv2d', 'Linear'], std=.02, bias=0.), dict(type='Constant', layer=['LayerNorm'], val=1., bias=0.),]
    my_module_weights_init(self.gt_seg_downsample_layers, my_module_init_cfg)

欢迎补充其他方便的方法。

相关推荐
菜鸟一枚在这2 分钟前
深度解析建造者模式:复杂对象构建的优雅之道
java·开发语言·算法
java1234_小锋6 分钟前
一周学会Flask3 Python Web开发-客户端状态信息Cookie以及加密
前端·python·flask·flask3
谨慎谦虚6 分钟前
Trae 体验:探索被忽视的 Chat 模式
人工智能·trae
北极的树8 分钟前
AI驱动的大前端开发工作流
人工智能
阿巴~阿巴~21 分钟前
多源 BFS 算法详解:从原理到实现,高效解决多源最短路问题
开发语言·数据结构·c++·算法·宽度优先
IT古董25 分钟前
【漫话机器学习系列】100.L2 范数(L2 Norm,欧几里得范数)
人工智能·机器学习
小天努力学java1 小时前
【面试系列】Java开发--AI常见面试题
java·人工智能·面试
B站计算机毕业设计超人1 小时前
计算机毕业设计Python+DeepSeek-R1高考推荐系统 高考分数线预测 大数据毕设(源码+LW文档+PPT+讲解)
大数据·python·机器学习·网络爬虫·课程设计·数据可视化·推荐算法
数造科技1 小时前
紧随“可信数据空间”政策风潮,数造科技正式加入开放数据空间联盟
大数据·人工智能·科技·安全·敏捷开发
@ V:ZwaitY091 小时前
如何打造TikTok矩阵:多账号管理与内容引流的高效策略
人工智能·矩阵·tiktok