Shortcuts

权重初始化

基于 Pytorch 构建模型时,我们通常会选择 nn.Module 作为模型的基类,搭配使用 Pytorch 的初始化模块 torch.nn.init,完成模型的初始化。MMEngine 在此基础上抽象出基础模块(BaseModule),让我们能够通过传参或配置文件来选择模型的初始化方式。此外,MMEngine 还提供了一系列模块初始化函数,让我们能够更加方便灵活地初始化模型参数。

配置式初始化

为了能够更加灵活地初始化模型权重,MMEngine 抽象出了模块基类 BaseModule。模块基类继承自 nn.Module,在具备 nn.Module 基础功能的同时,还支持在构造时接受参数,以此来选择权重初始化方式。继承自 BaseModule 的模型可以在实例化阶段接受 init_cfg 参数,我们可以通过配置 init_cfg 为模型中任意组件灵活地选择初始化方式。目前我们可以在 init_cfg 中配置以下初始化器:

Initializer Registered name Function
ConstantInit Constant 将 weight 和 bias 初始化为指定常量,通常用于初始化卷积
XavierInit Xavier 将 weight Xavier 方式初始化,将 bias 初始化成指定常量,通常用于初始化卷积
NormalInit Normal 将 weight 以正态分布的方式初始化,将 bias 初始化成指定常量,通常用于初始化卷积
TruncNormalInit TruncNormal 将 weight 以被截断的正态分布的方式初始化,参数 a 和 b 为正态分布的有效区域;将 bias 初始化成指定常量,通常用于初始化 Transformer
UniformInit Uniform 将 weight 以均匀分布的方式初始化,参数 a 和 b 为均匀分布的范围;将 bias 初始化为指定常量,通常用于初始化卷积
KaimingInit Kaiming 将 weight 以 Kaiming 的方式初始化,将 bias 初始化成指定常量,通常用于初始化卷积
Caffe2XavierInit Caffe2Xavier Caffe2 中 Xavier 初始化方式,在 Pytorch 中对应 "fan_in", "normal" 模式的 Kaiming 初始化,,通常用于初始化卷
Pretrained PretrainedInit 加载预训练权重

我们通过几个例子来理解如何在 init_cfg 里配置初始化器,来选择模型的初始化方式。

使用预训练权重初始化

假设我们定义了模型类 ToyNet,它继承自模块基类(BaseModule)。此时我们可以在 ToyNet 初始化时传入 init_cfg 参数来选择模型的初始化方式,实例化后再调用 init_weights 方法,完成权重的初始化。以加载预训练权重为例:

import torch
import torch.nn as nn

from mmengine.model import BaseModule


class ToyNet(BaseModule):

    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.conv1 = nn.Linear(1, 1)


# 保存预训练权重
toy_net = ToyNet()
torch.save(toy_net.state_dict(), './pretrained.pth')
pretrained = './pretrained.pth'

# 配置加载预训练权重的初始化方式
toy_net = ToyNet(init_cfg=dict(type='Pretrained', checkpoint=pretrained))
# 加载权重
toy_net.init_weights()
08/19 16:50:24 - mmengine - INFO - load model from: ./pretrained.pth
08/19 16:50:24 - mmengine - INFO - local loads checkpoint from path: ./pretrained.pth

init_cfg 是一个字典时,type 字段就表示一种初始化器,它需要被注册到 WEIGHT_INITIALIZERS 注册器。我们可以通过指定 init_cfg=dict(type='Pretrained', checkpoint='path/to/ckpt') 来加载预训练权重,其中 PretrainedPretrainedInit 初始化器的缩写,这个映射名由 WEIGHT_INITIALIZERS 维护;checkpointPretrainedInit 的初始化参数,用于指定权重的加载路径,它可以是本地磁盘路径,也可以是 URL。

备注

在所有的初始化器中,PretrainedInit 拥有最高的优先级。init_cfg 中其他初始化器初始化的权重会被 PretrainedInit 加载的预训练权重覆盖。

常用的初始化方式

和使用 PretrainedInit 初始化器类似,如果我们想对卷积做 Kaiming 初始化,需要令 init_cfg=dict(type='Kaiming', layer='Conv2d')。这样模型初始化时,就会以 Kaiming 初始化的方式来初始化类型为 Conv2d 的模块。

有时候我们可能需要用不同的初始化方式去初始化不同类型的模块,例如对卷积使用 Kaiming 初始化,对线性层使用 Xavier 初始化。此时我们可以使 init_cfg 成为一个列表,其中的每一个元素都表示对某些层使用特定的初始化方式。

import torch.nn as nn

from mmengine.model import BaseModule


class ToyNet(BaseModule):

    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.linear = nn.Linear(1, 1)
        self.conv = nn.Conv2d(1, 1, 1)


# 对卷积做 Kaiming 初始化,线性层做 Xavier 初始化
toy_net = ToyNet(
    init_cfg=[
        dict(type='Kaiming', layer='Conv2d'),
        dict(type='Xavier', layer='Linear')
    ], )
toy_net.init_weights()
08/19 16:50:24 - mmengine - INFO -
linear.weight - torch.Size([1, 1]):
XavierInit: gain=1, distribution=normal, bias=0

08/19 16:50:24 - mmengine - INFO -
linear.bias - torch.Size([1]):
XavierInit: gain=1, distribution=normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv.weight - torch.Size([1, 1, 1, 1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv.bias - torch.Size([1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

类似地,layer 参数也可以是一个列表,表示列表中的多种不同的 layer 均使用 type 指定的初始化方式

# 对卷积和线性层做 Kaiming 初始化
toy_net = ToyNet(init_cfg=[dict(type='Kaiming', layer=['Conv2d', 'Linear'])], )
toy_net.init_weights()
08/19 16:50:24 - mmengine - INFO -
linear.weight - torch.Size([1, 1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
linear.bias - torch.Size([1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv.weight - torch.Size([1, 1, 1, 1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv.bias - torch.Size([1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

更细粒度的初始化

有时同一类型的不同模块有不同初始化方式,例如现在有 conv1conv2 两个模块,他们的类型均为 Conv2d 。我们需要对 conv1 进行 Kaiming 初始化,conv2 进行 Xavier 初始化,则可以通过配置 override 参数来满足这样的需求:

import torch.nn as nn

from mmengine.model import BaseModule


class ToyNet(BaseModule):

    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.conv1 = nn.Conv2d(1, 1, 1)
        self.conv2 = nn.Conv2d(1, 1, 1)


# 对 conv1 做 Kaiming 初始化,对 从 conv2 做 Xavier 初始化
toy_net = ToyNet(
    init_cfg=[
        dict(
            type='Kaiming',
            layer=['Conv2d'],
            override=dict(name='conv2', type='Xavier')),
    ], )
toy_net.init_weights()
08/19 16:50:24 - mmengine - INFO -
conv1.weight - torch.Size([1, 1, 1, 1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv1.bias - torch.Size([1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv2.weight - torch.Size([1, 1, 1, 1]):
XavierInit: gain=1, distribution=normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv2.bias - torch.Size([1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

override 可以理解成一个嵌套的 init_cfg, 他同样可以是 list 或者 dict,也需要通过 type 字段指定初始化方式。不同的是 override 必须指定 namename 相当于 override 的作用域,如上例中,override 的作用域为 toy_net.conv2,我们会以 Xavier 初始化方式初始化 toy_net.conv2 下的所有参数,而不会影响作用域以外的模块。

自定义的初始化方式

尽管 init_cfg 能够控制各个模块的初始化方式,但是在不扩展 WEIGHT_INITIALIZERS 的情况下,我们是无法初始化一些自定义模块的,例如表格中提到的大多数初始化器,都需要对应的模块有 weightbias 属性 。对于这种情况,我们建议让自定义模块实现 init_weights 方法。模型调用 init_weights 时,会链式地调用所有子模块的 init_weights

假设我们定义了以下模块:

  • 继承自 nn.ModuleToyConv,实现了 init_weights 方法,让 custom_weight 初始化为 1,custom_bias 初始化为 0

  • 继承自模块基类的模型 ToyNet,且含有 ToyConv 子模块

我们在调用 ToyNetinit_weights 方法时,会链式的调用的子模块 ToyConvinit_weights 方法,实现自定义模块的初始化。

import torch
import torch.nn as nn

from mmengine.model import BaseModule


class ToyConv(nn.Module):

    def __init__(self):
        super().__init__()
        self.custom_weight = nn.Parameter(torch.empty(1, 1, 1, 1))
        self.custom_bias = nn.Parameter(torch.empty(1))

    def init_weights(self):
        with torch.no_grad():
            self.custom_weight = self.custom_weight.fill_(1)
            self.custom_bias = self.custom_bias.fill_(0)


class ToyNet(BaseModule):

    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.conv1 = nn.Conv2d(1, 1, 1)
        self.conv2 = nn.Conv2d(1, 1, 1)
        self.custom_conv = ToyConv()


toy_net = ToyNet(
    init_cfg=[
        dict(
            type='Kaiming',
            layer=['Conv2d'],
            override=dict(name='conv2', type='Xavier'))
    ])
# 链式调用 `ToyConv.init_weights()`,以自定义的方式初始化
toy_net.init_weights()
08/19 16:50:24 - mmengine - INFO -
conv1.weight - torch.Size([1, 1, 1, 1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv1.bias - torch.Size([1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv2.weight - torch.Size([1, 1, 1, 1]):
XavierInit: gain=1, distribution=normal, bias=0

08/19 16:50:24 - mmengine - INFO -
conv2.bias - torch.Size([1]):
KaimingInit: a=0, mode=fan_out, nonlinearity=relu, distribution =normal, bias=0

08/19 16:50:24 - mmengine - INFO -
custom_conv.custom_weight - torch.Size([1, 1, 1, 1]):
Initialized by user-defined `init_weights` in ToyConv

08/19 16:50:24 - mmengine - INFO -
custom_conv.custom_bias - torch.Size([1]):
Initialized by user-defined `init_weights` in ToyConv

小结

最后我们对 init_cfginit_weights 两种初始化方式做一些总结:

1. 配置 init_cfg 控制初始化

  • 通常用于初始化一些比较底层的模块,例如卷积、线性层等。如果想通过 init_cfg 配置自定义模块的初始化方式,需要将相应的初始化器注册到 WEIGHT_INITIALIZERS 里。

  • 动态初始化特性,初始化方式随 init_cfg 的值改变。

2. 实现 init_weights 方法

  • 通常用于初始化自定义模块。相比于 init_cfg 的自定义初始化,实现 init_weights 方法更加简单,无需注册。然而,它的灵活性不及 init_cfg,无法动态地指定任意模块的初始化方式。

备注

  • init_weights 的优先级比 init_cfg

  • 执行器会在 train() 函数中调用 init_weights。

函数式初始化

自定义的初始化方式一节提到,我们可以在 init_weights 里实现自定义的参数初始化逻辑。为了能够更加方便地实现参数初始化,MMEngine 在 torch.nn.init的基础上,提供了一系列模块初始化函数来初始化整个模块。例如我们对卷积层的权重(weight)进行正态分布初始化,卷积层的偏置(bias)进行常数初始化,基于 torch.nn.init 的实现如下:

from torch.nn.init import normal_, constant_
import torch.nn as nn

model = nn.Conv2d(1, 1, 1)
normal_(model.weight, mean=0, std=0.01)
constant_(model.bias, val=0)
Parameter containing:
tensor([0.], requires_grad=True)

上述流程实际上是卷积正态分布初始化的标准流程,因此 MMEngine 在此基础上做了进一步地简化,实现了一系列常用的模块初始化函数。相比 torch.nn.init,MMEngine 提供的初始化函数直接接受卷积模块,一行代码能实现同样的初始化逻辑:

from mmengine.model import normal_init

normal_init(model, mean=0, std=0.01, bias=0)

类似地,我们也可以用 Kaiming 初始化和 Xavier 初始化:

from mmengine.model import kaiming_init, xavier_init

kaiming_init(model)
xavier_init(model)

目前 MMEngine 提供了以下初始化函数:

初始化函数 功能
constant_init 将 weight 和 bias 初始化为指定常量,通常用于初始化卷积
xavier_init 将 weight 以 Xavier 方式初始化,将 bias 初始化成指定常量,通常用于初始化卷积
normal_init 将 weight 以正态分布的方式初始化,将 bias 初始化成指定常量,通常用于初始化卷积
trunc_normal_init 将 weight 以被截断的正态分布的方式初始化,参数 a 和 b 为正态分布的有效区域;将 bias 初始化成指定常量,通常用于初始化 Transformer
uniform_init 将 weight 以均匀分布的方式初始化,参数 a 和 b 为均匀分布的范围;将 bias 初始化为指定常量,通常用于初始化卷积
kaiming_init 将 weight 以 Kaiming 方式初始化,将 bias 初始化成指定常量,通常用于初始化卷积
caffe2_xavier_init Caffe2 中 Xavier 初始化方式,在 Pytorch 中对应 "fan_in", "normal" 模式的 Kaiming 初始化,通常用于初始化卷积
bias_init_with_prob 以概率值的形式初始化 bias
Read the Docs v: stable
Versions
latest
stable
v0.10.2
v0.10.1
v0.10.0
v0.9.1
v0.9.0
v0.8.5
v0.8.4
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.