本章程式碼：https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson4/grad_vanish_explod.py

在搭建好網路模型之後，一個重要的步驟就是對網路模型中的權值進行初始化。適當的權值初始化可以加快模型的收斂，而不恰當的權值初始化可能引發梯度消失或者梯度爆炸，最終導致模型無法收斂。下面分 3 部分介紹。第一部分介紹不恰當的權值初始化是如何引發梯度消失與梯度爆炸的，第二部分介紹常用的 Xavier 方法與 Kaiming 方法，第三部分介紹 PyTorch 中的 10 種初始化方法。

梯度消失與梯度爆炸

考慮一個 3 層的全連線網路。

$H_{1}=X \times W_{1}$，$H_{2}=H_{1} \times W_{2}$，$Out=H_{2} \times W_{3}$

$\begin{aligned} \Delta \mathrm{W}{2} &=\frac{\partial \mathrm{Loss}}{\partial \mathrm{W}{2}}=\frac{\partial \mathrm{Loss}}{\partial \mathrm{out}} * \frac{\partial \mathrm{out}}{\partial \mathrm{H}{2}} * \frac{\partial \mathrm{H}

所以 $\Delta \mathrm{W}{2}$ 依賴於前一層的輸出 $H{1}$。如果 $H_{1}$ 趨近於零，那麼 $\Delta \mathrm{W}{2}$ 也接近於 0，造成梯度消失。如果 $H{1}$ 趨近於無窮大，那麼 $\Delta \mathrm{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)

輸出為：

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>)

也就是資料太大(梯度爆炸)或者太小(梯度消失)了。接下來我們在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

輸出如下：

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

可以看到每一層的標準差是越來越大的，並在在 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)$：兩個相互獨立的隨機變數之和的方差等於它們的方差的和。

可以推匯出兩個隨機變數的乘積的方差如下：

$D(X \times Y)=E[(XY)^{2}] - [E(XY)]^{2}=D(X) \times D(Y) + D(X) \times [E(Y)]^{2} + D(Y) \times [E(X)]^{2}$

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

我們以輸入層第一個神經元為例：

$\mathrm{H}{11}=\sum{i=0}^{n} X_{i} \times W_{1 i}$

其中輸入 X 和權值 W 都是服從 $N(0,1)$ 的正態分佈，所以這個神經元的方差為：

$\begin{aligned} \mathbf{D}\left(\mathrm{H}{11}\right) &=\sum{i=0}^{n} \boldsymbol{D}\left(X_{i}\right) * \boldsymbol{D}\left(W_{1 i}\right) \ &=n *(1 * 1) \ &=n \end{aligned}$

標準差為：$\operatorname{std}\left(\mathrm{H}{11}\right)=\sqrt{\mathbf{D}\left(\mathrm{H}{11}\right)}=\sqrt{n}$，所以每經過一個網路層，方差就會擴大 n 倍，標準差就會擴大 $\sqrt{n}$ 倍，n 為每層神經元個數，直到超出數值表示範圍。對比上面的程式碼可以看到，每層神經元個數為 256，輸出資料的標準差為 1，所以第一個網路層輸出的標準差為 16 左右，第二個網路層輸出的標準差為 256 左右，以此類推，直到 31 層超出資料表示範圍。可以把每層神經元個數改為 400，那麼每層標準差擴大 20 倍左右。從 $D(\mathrm{H}{11})=\sum{i=0}^{n} D(X_{i}) \times D(W_{1 i})$，可以看出，每一層網路輸出的方差與神經元個數、輸入資料的方差、權值方差有關，其中比較好改變的是權值的方差 $D(W)$，所以 $D(W)= \frac{1}{n}$，標準差為 $std(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 左右。通過恰當的權值初始化，可以保持權值在更新過程中維持在一定範圍之內，不過過大，也不會過小。

上述是沒有使用非線性變換的實驗結果，如果在forward()中新增非線性變換tanh，每一層的輸出方差還是會越來越小，會導致梯度消失。因此出現了 Xavier 初始化方法與 Kaiming 初始化方法。

Xavier 方法與 Kaiming 方法

Xavier 方法

Xavier 是 2010 年提出的，針對有非線性啟用函式時的權值初始化方法，目標是保持資料的方差維持在 1 左右，主要針對飽和啟用函式如 sigmoid 和 tanh 等。同時考慮前向傳播和反向傳播，需要滿足兩個等式：$\boldsymbol{n}{\boldsymbol{i}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1}$ 和 $\boldsymbol{n}{\boldsymbol{i+1}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{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}$，解得：$\boldsymbol{a}=\frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}$，所以 $W$ 服從分佈 $U\left[-\frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}, \frac{\sqrt{6}}{\sqrt{n_{i}+n_{i+1}}}\right]$

所以初始化方法改為：

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)

並且每一層的啟用函式都使用 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

可以看到每層輸出的方差都維持在 0.6 左右。

PyTorch 也提供了 Xavier 初始化方法，可以直接呼叫：

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

nn.init.calculate_gain()

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

nn.init.calculate_gain(nonlinearity,param=**None**)的主要功能是經過一個分佈的方差經過啟用函式後的變化尺度，主要有兩個引數：

• nonlinearity：啟用函式名稱
• param：啟用函式的引數，如 Leaky ReLU 的 negative_slop。

下面是計算標準差經過啟用函式的變化尺度的程式碼。

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)

輸出如下：

gain:1.5982500314712524
tanh_gain in PyTorch: 1.6666666666666667

結果表示，原有資料分佈的方差經過 tanh 之後，標準差會變小 1.6 倍左右。

Kaiming 方法

雖然 Xavier 方法提出了針對飽和啟用函式的權值初始化方法，但是 AlexNet 出現後，大量網路開始使用非飽和的啟用函式如 ReLU 等，這時 Xavier 方法不再適用。2015 年針對 ReLU 及其變種等啟用函式提出了 Kaiming 初始化方法。

針對 ReLU，方差應該滿足：$\mathrm{D}(W)=\frac{2}{n_{i}}$；針對 ReLu 的變種，方差應該滿足：$\mathrm{D}(W)=\frac{2}{n_{i}}$，a 表示負半軸的斜率，如 PReLU 方法，標準差滿足 $\operatorname{std}(W)=\sqrt{\frac{2}{\left(1+a^{2}\right) * n_{i}}}$。程式碼如下：nn.init.normal_(m.weight.data, std=np.sqrt(2 / self.neural_num))，或者使用 PyTorch 提供的初始化方法：nn.init.kaiming_normal_(m.weight.data)，同時把啟用函式改為 ReLU。

常用初始化方法

PyTorch 中提供了 10 中初始化方法

1. Xavier 均勻分佈
2. Xavier 正態分佈
3. Kaiming 均勻分佈
4. Kaiming 正態分佈
5. 均勻分佈
6. 正態分佈
7. 常數分佈
8. 正交矩陣初始化
9. 單位矩陣初始化
10. 稀疏矩陣初始化

每種初始化方法都有它自己適用的場景，原則是保持每一層輸出的方差不能太大，也不能太小。

參考資料

• 深度之眼 PyTorch 框架班

如果你覺得這篇文章對你有幫助，不妨點個贊，讓我有更多動力寫出好文章。