【Pytorch】模型权重的初始化函数

在神经网络模型中，参数权重的初始设置非常重要，其合理设置能够保证模型的收敛以及数值运算的速度。

pytorch中常用的初始化函数封装在torch.nn.init下，常见策略主要包括：

1. 均匀分布初始化

"""
a: 均匀分布下限
b: 均匀分布上限
返回同tensor同shape的初始化张量
"""
init.uniform_(tensor, a=0, b=1)

2. 正态分布初始化

"""
mean: 正态分布均值
std: 正态分布方差
返回同tensor同shape的初始化张量
"""
init.normal_(tensor, mean=0, std=1)

3. 常量初始化

"""
val: 常数
返回同tensor同shape的初始化张量，张量元素均为val
"""
init.normal_(tensor, val=0)

4. Xavier初始化

一种基于模型参数初始化后，每层输出的方差不该受该层输入个数影响，且每层梯度的方差也不该受该层输出个数影响思想的初始化策略，特别适用于激活函数为tanh函数，能够有效减少梯度消失和梯度爆炸问题。pytorch中提供了uniform和normal两种策略：

"""
gain为缩放因子，uniform和normal的统计参数根据公式自动计算得到
"""
# uniform策略
init.xavier_uniform_(tensor, gain=1)

# normal策略
init.xavier_normal_(tensor, gain=1)

关于这两种策略的具体计算公式如下：

1）uniform策略
返回均匀分布$U～(-a,a)$的随机采样值，其中参数$a$根据下式计算：
$$a=gain\times \sqrt{\frac{6}{fa{n}_{in}+fa{n}_{out}}}$$其中参数$fa{n}_{in},fa{n}_{out}$根据输入的tensor自动计算：
( a ) tensor的维度必须不小于2；

( b ) 若tensor的维度为2（如线性回归模型的特征权重），则$fa{n}_{in}$为tensor的第二个维度（即输出特征个数），而$fa{n}_{out}$为tensor的第一个维度（即输入特征个数）。

( c ) 若tensor的维度大于2（如卷积核），则$fa{n}_{in}$为output_channel个数×每个卷积核的参数个数，$fa{n}_{out}$为in_channel个数×每个卷积核的参数个数。

其对应源码见下：

def _calculate_fan_in_and_fan_out(tensor):
    dimensions = tensor.dim()
    if dimensions < 2:
        raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions")

    if dimensions == 2:  # Linear
        fan_in = tensor.size(1)
        fan_out = tensor.size(0)
    else:
        num_input_fmaps = tensor.size(1)
        num_output_fmaps = tensor.size(0)
        receptive_field_size = 1
        if tensor.dim() > 2:
            receptive_field_size = tensor[0][0].numel()
        fan_in = num_input_fmaps * receptive_field_size
        fan_out = num_output_fmaps * receptive_field_size

    return fan_in, fan_out

2）normal策略
返回正态分布$N～(0,std)$的随机采样值，其中参数$std$根据下式计算：
$$std=gain\times \sqrt{\frac{2}{fa{n}_{in}+fa{n}_{out}}}$$其中参数$fa{n}_{in},fa{n}_{out}$的计算同上。

5. He初始化

该初始化方法由何凯明提出，特别适用于Relu激活函数及其变种，其基本思想与Xavier类似。pytorch中同样提供了uniform和normal两种策略：

"""
a: 激活函数在负轴区域内的斜率，对于Relu，该值默认值为0
mode: 可设为'fan_in'或'fan_out'，分别表示控制前向传播或后向传播过程中输入、输出张量的方差一致，
在计算随机分布参数时，也会使用赌赢的参数。默认'fan_in'。
nonlinearity：可配置各种激活函数，但一般只建议采用'leaky_relu'或'relu'
"""
# uniform策略
init.kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')

# uniform策略
init.kaiming_uniform_(tensor, a=0, mode='fan_in'; nonlinearity='leaky_relu')

这两种计算策略的具体分布参数的计算均取决于mode和nonlinearity参数的设置，具体可参见源码。这里仅给出mode='fan_in',nonlinearity='leaky_relu'时的计算公式。

1）uniform策略
返回均匀分布$U～(-b,b)$的随机采样值，其中参数$b$根据下式计算：$$b=\sqrt{\frac{6}{(1+a^2)\times fa{n}_{in}}}$$其中参数$fa{n}_{in}$的计算同上文。

pytorch中，卷积网络的默认初始化就采用了这种策略，具体可见源码中_ConvNd类的reset_parameters函数。

2）normal策略
返回均匀分布$N～(0,std)$的随机采样值，其中参数$b$根据下式计算：$$b=\sqrt{\frac{2}{(1+a^2)\times fa{n}_{in}}}$$其中参数$fa{n}_{in}$的计算同上文。

【Reference】

1. pytorch官方文档
2. pytorch nn.init 中实现的初始化函数