权值初始化

梯度消失与梯度爆炸

考虑一个 3 层的全连接网络。

$H 1 = X \times W 1 ， H 2 = H 1 \times W 2 ， O u t = H 2 \times W_{3}$

Untitled

\begin{aligned} Δ W 2 & = \frac{\partial Loss}{\partial W 2} = \frac{\partial Loss}{\partial out} \frac{\partial out}{\partial H_{2}} \frac{\partial H 2}{\partial w 2} = \frac{\partial Loss}{\partial out} \frac{\partial out}{\partial H_{2}} H_{1} \end{aligned}

所以 $Δ W 2$ 依赖于前一层的输出 $H 1$ 。如果 $H 1$ 趋近于零，那么 $Δ W 2$ 也接近于 0，造成梯度消失。如果 $H 1$ 趋近于无穷大，那么 $Δ W 2$ 也接近于无穷大，造成梯度爆炸。要避免梯度爆炸或者梯度消失，就要严格控制网络层输出的数值范围。

下面构建 100 层全连接网络，先不适用非线性激活函数，每层的权重初始化为服从 $N (0, 1)$ 的正态分布，输出数据使用随机初始化的数据。

import torch
import torch.nn as nn
from common_tools import set_seed

set_seed(1)  # 设置随机种子

class MLP(nn.Module):
    def __init__(self, neural_num, layers):
        super(MLP, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(neural_num, neural_num, bias=False) for i in range(layers)])
        self.neural_num = neural_num

    def forward(self, x):
        for (i, linear) in enumerate(self.linears):
            x = linear(x)
        return x

    def initialize(self):
        for m in self.modules():
            # 判断这一层是否为线性层，如果为线性层则初始化权值
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight.data)    # normal: mean=0, std=1

layer_nums = 100
neural_nums = 256
batch_size = 16

net = MLP(neural_nums, layer_nums)
net.initialize()

inputs = torch.randn((batch_size, neural_nums))  # normal: mean=0, std=1

output = net(inputs)
print(output)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

输出为：

tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=<MmBackward>)

1
2
3
4
5
6
7

也就是数据太大(梯度爆炸)或者太小(梯度消失)了。接下来我们在forward()函数中判断每一次前向传播的输出的标准差是否为 nan，如果是 nan 则停止前向传播。

def forward(self, x):
        for (i, linear) in enumerate(self.linears):
            x = linear(x)

            print("layer:{}, std:{}".format(i, x.std()))
            if torch.isnan(x.std()):
                print("output is nan in {} layers".format(i))
                break

        return x

1
2
3
4
5
6
7
8
9
10

输出如下：

layer:0, std:15.959932327270508
layer:1, std:256.6237487792969
layer:2, std:4107.24560546875
.
.
.
layer:29, std:1.322983152787379e+36
layer:30, std:2.0786820453988485e+37
layer:31, std:nan
output is nan in 31 layers

1
2
3
4
5
6
7
8
9
10

可以看到每一层的标准差是越来越大的，并在在 31 层时超出了数据可以表示的范围。

下面推导为什么网络层输出的标准差越来越大。

首先给出 3 个公式：

$E (X \times Y) = E (X) \times E (Y)$ ：两个相互独立的随机变量的乘积的期望等于它们的期望的乘积。

$D (X) = E (X^{2}) - [E (X)]^{2}$ ：一个随机变量的方差等于它的平方的期望减去期望的平方

$D (X + Y) = D (X) + D (Y)$ ：两个相互独立的随机变量之和的方差等于它们的方差的和。

可以推导出两个随机变量的乘积的方差如下：

如 果 E (X) = 0 ， E (Y) = 0 ， 那 么 D (X \times Y) = D (X) \times D (Y)

我们以输入层第一个神经元为例：

H_{11} = \sum_{i = 0}^{n} X_{i} \times W_{1 i}

其中输入 $X$ 和权值 $W$ 都是服从 $N (0, 1)$ 的正态分布，所以这个神经元的方差为：

\begin{aligned} D (H_{11}) & = \sum_{i = 0}^{n} D (X_{i}) * D (W_{1 i}) = n * (1 * 1) = n \end{aligned}

标准差为： $std (H_{11}) = \sqrt{D (H_{11})} = \sqrt{n}$ ，所以每经过一个网络层，方差就会扩大 n 倍，标准差就会扩大 $\sqrt{n}$ 倍， $n$ 为每层神经元个数，直到超出数值表示范围。对比上面的代码可以看到，每层神经元个数为 256，输出数据的标准差为 1，所以第一个网络层输出的标准差为 16 左右，第二个网络层输出的标准差为 256 左右，以此类推，直到 31 层超出数据表示范围。可以把每层神经元个数改为 400，那么每层标准差扩大 20 倍左右。从 $D (H_{11}) = \sum_{i = 0}^{n} D (X_{i}) \times D (W_{1 i})$ ，可以看出，每一层网络输出的方差与神经元个数、输入数据的方差、权值方差有关，其中比较好改变的是权值的方差 $D (W)$ ，所以 $D (W) = \frac{1}{n}$ ，标准差为 $s t d (W) = \sqrt{\frac{1}{n}}$ 。因此修改权值初始化代码为nn.init.normal_(m.weight.data, std=np.sqrt(1/self.neural_num)),结果如下：

layer:0, std:0.9974957704544067
layer:1, std:1.0024365186691284
layer:2, std:1.002745509147644
.
.
.
layer:94, std:1.031973123550415
layer:95, std:1.0413124561309814
layer:96, std:1.0817031860351562

1
2
3
4
5
6
7
8
9

修改之后，没有出现梯度消失或者梯度爆炸的情况，每层神经元输出的方差均在 1 左右。通过恰当的权值初始化，可以保持权值在更新过程中维持在一定范围之内，不会过大，也不会过小。

上述是没有使用非线性变换的实验结果，如果在forward()中添加非线性变换tanh，每一层的输出方差还是会越来越小，会导致梯度消失。因此出现了 Xavier 初始化方法与 Kaiming 初始化方法。

torch.nn.init

torch.nn.init 是 PyTorch 中的一个模块，它提供了多种权重初始化方法。权重初始化是神经网络训练过程中的一个重要步骤，合适的初始化方法可以帮助模型更好地收敛，提高训练速度和性能。

torch.nn.init 提供了以下一些常用的初始化方法：

常量初始化

nn.init.constant_: 将张量填充为给定的常量值。

基于分布的初始化

nn.init.uniform_: 从均匀分布中抽取样本并用它们填充张量。
nn.init.normal_: 从正态分布（高斯分布）中抽取样本并用它们填充张量。
nn.init.xavier_uniform_: 使用 Glorot 初始化（也称作 Xavier 初始化），从均匀分布中抽取样本，并根据输入和输出单元数的数量来调整这些样本的范围。
nn.init.xavier_normal_: 使用 Glorot 初始化从正态分布中抽取样本。
nn.init.kaiming_uniform_: 使用 He 初始化（也称作 Kaiming 初始化）从均匀分布中抽取样本，它特别适用于 ReLU 激活函数。
nn.init.kaiming_normal_: 使用 He 初始化从正态分布中抽取样本。

稀疏初始化

nn.init.orthogonal_: 使用正交矩阵填充张量。
nn.init.sparse_: 用稀疏矩阵填充张量。

其他初始化

nn.init.eye_: 用单位矩阵填充张量。
nn.init.dirac_: 在特定维度上创建一个 "Dirac" delta 分布。
nn.init.calculate_gain: 计算用于初始化方法的缩放因子。

除了calculate_gain，所有函数的后缀都带有下划线，意味着这些函数将会直接原地更改输入张量的值

使用示例

我们通常会根据实际模型来使用torch.nn.init进行初始化，通常使用isinstance()来进行判断模块属于什么类型。

import torch
import torch.nn as nn

conv = nn.Conv2d(1,3,3)
linear = nn.Linear(10,1)

print(isinstance(conv,nn.Conv2d)) # 判断conv是否是nn.Conv2d类型
print(isinstance(linear,nn.Conv2d)) # 判断linear是否是nn.Conv2d类型

1
2
3
4
5
6
7
8

True
False

1
2

# 查看随机初始化的conv参数
conv.weight.data
# 查看linear的参数
linear.weight.data

1
2
3
4

tensor([[[[ 0.1174,  0.1071,  0.2977],
          [-0.2634, -0.0583, -0.2465],
          [ 0.1726, -0.0452, -0.2354]]],
        [[[ 0.1382,  0.1853, -0.1515],
          [ 0.0561,  0.2798, -0.2488],
          [-0.1288,  0.0031,  0.2826]]],
        [[[ 0.2655,  0.2566, -0.1276],
          [ 0.1905, -0.1308,  0.2933],
          [ 0.0557, -0.1880,  0.0669]]]])

tensor([[-0.0089,  0.1186,  0.1213, -0.2569,  0.1381,  0.3125,  0.1118, -0.0063, -0.2330,  0.1956]])

1
2
3
4
5
6
7
8
9
10
11

对于不同的类型层，我们就可以设置不同的权值初始化的方法。

# 对conv进行kaiming初始化
torch.nn.init.kaiming_normal_(conv.weight.data)
conv.weight.data
# 对linear进行常数初始化
torch.nn.init.constant_(linear.weight.data,0.3)
linear.weight.data

1
2
3
4
5
6

tensor([[[[ 0.3249, -0.0500,  0.6703],
          [-0.3561,  0.0946,  0.4380],
          [-0.9426,  0.9116,  0.4374]]],
        [[[ 0.6727,  0.9885,  0.1635],
          [ 0.7218, -1.2841, -0.2970],
          [-0.9128, -0.1134, -0.3846]]],
        [[[ 0.2018,  0.4668, -0.0937],
          [-0.2701, -0.3073,  0.6686],
          [-0.3269, -0.0094,  0.3246]]]])
tensor([[0.3000, 0.3000, 0.3000, 0.3000, 0.3000, 0.3000, 0.3000, 0.3000, 0.3000,0.3000]])

1
2
3
4
5
6
7
8
9
10

Xavier 方法

Xavier 是 2010 年提出的，针对有非线性激活函数时的权值初始化方法，目标是保持数据的方差维持在 1 左右，主要针对饱和激活函数如 sigmoid 和 tanh 等。同时考虑前向传播和反向传播，需要满足两个等式： $n_{i} * D (W) = 1$ 和 $n_{i + 1} * D (W) = 1$ 可得： $D (W) = \frac{2}{n_{i} + n_{i + 1}}$ 。为了使 Xavier 方法初始化的权值服从均匀分布，假设 $W$ 服从均匀分布 $U [- a, a]$ ，那么方差 $D (W) = \frac{(- a - a)^{2}}{12} = \frac{(2 a)^{2}}{12} = \frac{a^{2}}{3}$ ，令 $\frac{2}{n_{i} + n_{i + 1}} = \frac{a^{2}}{3}$ ，解得： $a = \frac{\sqrt{6}}{\sqrt{n_{i} + n_{i + 1}}}$ ，所以 $W$ 服从分布 $U [- \frac{\sqrt{6}}{\sqrt{n_{i} + n_{i + 1}}}, \frac{\sqrt{6}}{\sqrt{n_{i} + n_{i + 1}}}]$

所以初始化方法改为：

a = np.sqrt(6 / (self.neural_num + self.neural_num))
# 把 a 变换到 tanh，计算增益
tanh_gain = nn.init.calculate_gain('tanh')
a *= tanh_gain

nn.init.uniform_(m.weight.data, -a, a)

1
2
3
4
5
6

并且每一层的激活函数都使用 tanh，输出如下：

layer:0, std:0.7571136355400085
layer:1, std:0.6924336552619934
layer:2, std:0.6677976846694946
.
.
.
layer:97, std:0.6426210403442383
layer:98, std:0.6407480835914612
layer:99, std:0.6442216038703918

1
2
3
4
5
6
7
8
9

可以看到每层输出的方差都维持在 0.6 左右。 PyTorch 也提供了 Xavier 初始化方法，可以直接调用：

tanh_gain = nn.init.calculate_gain('tanh')
nn.init.xavier_uniform_(m.weight.data, gain=tanh_gain)

1
2

nn.init.calculate_gain()

上面的初始化方法都使用了tanh_gain = nn.init.calculate_gain('tanh')。

- nonlinearity：激活函数名称
- param：激活函数的参数，如 Leaky ReLU 的 negative_slop。

1
2

下面是计算标准差经过激活函数的变化尺度的代码。

x = torch.randn(10000) out = torch.tanh(x)

gain = x.std() / out.std() print('gain:{}'.format(gain))

tanh_gain = nn.init.calculate_gain('tanh') print('tanh_gain in PyTorch:', tanh_gain)

1
2
3
4
5

输出如下：

gain:1.5982500314712524 tanh_gain in PyTorch: 1.6666666666666667

1
2

结果表示，原有数据分布的方差经过 tanh 之后，标准差会变小 1.6倍左右。

Kaiming 方法

虽然 Xavier 方法提出了针对饱和激活函数的权值初始化方法，但是 AlexNet 出现后，大量网络开始使用非饱和的激活函数如 ReLU 等，这时 Xavier 方法不再适用。2015 年针对 ReLU 及其变种等激活函数提出了 Kaiming 初始化方法。

针对 ReLU，方差应该满足： $D (W) = \frac{2}{n_{i}}$ ；针对 ReLu 的变种，方差应该满足： $D (W) = \frac{2}{n_{i}}$ ， $a$ 表示负半轴的斜率，如 PReLU 方法，标准差满足 $std (W) = \sqrt{\frac{2}{(1 + a^{2}) * n_{i}}}$ 。

#PyTorch

上次更新: 2025/04/02, 12:03:38

← 模型参数模型保存与加载→