Pytorch 模型初始化

在开始阅读下面的内容之前，我们需要牢记参数初始化的目的是为了让神经网络在训练过程中学习到有用的信息，这意味着参数梯度不应该为0。而我们知道在全连接的神经网络中，参数梯度和反向传播得到的状态梯度以及入激活值有关——激活值饱和会导致该层状态梯度信息为0，然后导致下面所有层的参数梯度为0；入激活值为0会导致对应参数梯度为0。所以如果要保证参数梯度不等于0，那么参数初始化应该使得各层激活值不会出现饱和现象且激活值不为0。我们把这两个条件总结为参数初始化条件：

初始化必要条件一：各层激活值不会出现饱和现象。
初始化必要条件二：各层激活值不为0。

1. 常见初始化方式

   参数初始化分为：固定值初始化、预训练初始化、随机初始化。

   固定值初始化

   固定初始化是指将模型参数初始化为一个固定的常数，这意味着所有单元具有相同的初始化状态，所有的神经元都具有相同的输出和更新梯度，并进行完全相同的更新，这种初始化方法使得神经元间不存在非对称性，从而使得模型效果大打折扣。

   预训练初始化

   预训练初始化是神经网络初始化的有效方式，比较早期的方法是使用 greedy layerwise auto-encoder 做无监督学习的预训练，经典代表为 Deep Belief Network；而现在更为常见的是有监督的预训练+模型微调。

   随机初始化

   随机初始化是指随机进行参数初始化，但如果不考虑随机初始化的分布则会导致梯度爆炸和梯度消失的问题。

2. 随机初始化介绍

   先介绍三个用的比较多的初始化方法：高斯分布、均匀分布、常数分布，以及全0分布、全1分布

#高斯分布
torch.nn.init.normal_(tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0) → torch.Tensor
#均匀分布
torch.nn.init.uniform_(tensor: torch.Tensor, a: float = 0.0, b: float = 1.0) → torch.Tensor
#常数分布
torch.nn.init.constant_(tensor: torch.Tensor, val: float) → torch.Tensor
#全0分布
torch.nn.init.zeros_(tensor: torch.Tensor) → torch.Tensor
#全1分布
torch.nn.init.ones_(tensor: torch.Tensor) → torch.Tensor

3. 两种常见的模型初始化方法

   Xavier初始化方法

   Xavier初始化:

   条件：正向传播时，激活值的方差保持不变；反向传播时，关于状态值的梯度的方差保持不变。

   初始化方法：
   
   需要满足激活函数关于0对称，且主要针对于全连接神经网络。适用于tanh和softsign

   Xavier初始化的缺点

   这并不适用于ReLU激活函数。另一个是激活值关于0对称，这个不适用于sigmoid函数和ReLU函数。

   源码

#均匀分布
torch.nn.init.xavier_uniform_(tensor: torch.Tensor, gain: float = 1.0) → torch.Tensor
#正态分布
torch.nn.init.xavier_normal_(tensor: torch.Tensor, gain: float = 1.0) → torch.Tensor

Kaiming初始化

条件：正向传播时，状态值的方差保持不变；反向传播时，关于激活值的梯度的方差保持不变。

源码

#均匀分布
torch.nn.init.kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')
#正态分布
torch.nn.init.kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')

4. Pytorch 默认初始化方式

nn.Conv2d 、 nn.ConvTranspose2d 、nn.Linear

    def reset_parameters(self):
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)

nn.BatchNorm2d

    def reset_running_stats(self):
        if self.track_running_stats:
            self.running_mean.zero_()
            self.running_var.fill_(1)
            self.num_batches_tracked.zero_()

    def reset_parameters(self):
        self.reset_running_stats()
        if self.affine:
            init.ones_(self.weight)
            init.zeros_(self.bias)

5. 自定义的初始化方式 以DenseNet为例

# params initialization
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                     nn.init.constant_(m.bias, 0)