快捷導(dǎo)航

PyTorch權(quán)值初始化原理解析

更新時(shí)間：2023年07月17日 09:07:45 作者：YOLO

這篇文章主要為大家介紹了PyTorch權(quán)值初始化原理示例解析，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

Pytorch：權(quán)值初始化

在搭建好網(wǎng)絡(luò)模型之后，首先需要對(duì)網(wǎng)絡(luò)模型中的權(quán)值進(jìn)行初始化。權(quán)值初始化的作用有很多，通常，一個(gè)好的權(quán)值初始化將會(huì)加快模型的收斂，而比較差的權(quán)值初始化將會(huì)引發(fā)梯度爆炸或者梯度消失。下面將具體解釋其中的原因：

梯度消失與梯度爆炸

考慮一個(gè) 3 層的全連接網(wǎng)絡(luò)。

$H_{1}=X \times W_{1}$，$H_{2}=H_{1} \times W_{2}$，$Out=H_{2} \times W_{3}$，如下圖所示，

其中第 2 層的權(quán)重梯度如下：

$\begin{array}{l}
\mathrm{H}_{2}=\mathrm{H}_{1} * \mathrm{~W}_{\mathbf{2}} \
\Delta \mathrm{W}_{\mathbf{2}}=\frac{\partial \text { Loss }}{\partial \mathrm{W}_{2}}=\frac{\partial \mathrm{Loss}}{\partial \text { out }} \frac{\partial \text { out }}{\partial \mathrm{H}_{2}} \frac{\partial \mathrm{H}_{2}}{\partial \mathrm{W}_{2}} \ =\frac{\partial \text { Loss }}{\partial \text { out }} \frac{\partial \text { out }}{\partial \mathrm{H}_{2}} \mathrm{H}_{1} \
\end{array}$

由上式化簡可知，如果H\_1發(fā)生以下變化，那么對(duì)應(yīng)的梯度也就會(huì)發(fā)生變化：

梯度消失: $\mathrm{H}_{1} \rightarrow 0 \Rightarrow \Delta \mathrm{W}_{2} \rightarrow 0$

梯度爆炸: $\mathrm{H}_{1} \rightarrow \infty \Rightarrow \Delta \mathrm{W}_{2} \rightarrow \infty $

因此，為了避免以上兩種情況，就必須嚴(yán)格控制網(wǎng)絡(luò)層輸出的數(shù)值范圍。

具體可以通過構(gòu)建 100 層全連接網(wǎng)絡(luò)，先不使用非線性激活函數(shù)，每層的權(quán)重初始化為服從 $N(0,1)$ 的正態(tài)分布，輸出數(shù)據(jù)使用隨機(jī)初始化的數(shù)據(jù)，這樣的例子來直觀地感受影響：

import torch
import torch.nn as nn
from common_tools import set_seed
set_seed(1)  # 設(shè)置隨機(jī)種子
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():
            # 判斷這一層是否為線性層，如果為線性層則初始化權(quán)值
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight.data)    # normal: mean=0, std=1
# 網(wǎng)絡(luò)的層數(shù)
layer_nums = 100
# 神經(jīng)元的個(gè)數(shù)
neural_nums = 256
batch_size = 16
net = MLP(neural_nums, layer_nums)
net.initialize()
# 設(shè)置隨機(jī)初始化的輸入
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>)

通過輸出可知，輸出值均為nan，即非數(shù)字類型，原因可能是數(shù)據(jù)太大(梯度爆炸)或者太小(梯度消失)。

為了具體知道是在哪一層開始出現(xiàn)nan的，我們可以在forward函數(shù)中添加判斷得知，查看每一次前向轉(zhuǎn)播的標(biāo)準(zhǔn)差是否是nan，若是，則停止前向傳播并輸出。

這里判斷是否為nan時(shí)采用了 torch.isnan 函數(shù)

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

可見，之際上輸出的標(biāo)準(zhǔn)差是逐層遞增的，具體為什么會(huì)導(dǎo)致這種情況：

$E(X \times Y)=E(X) \times E(Y)$：兩個(gè)相互獨(dú)立的隨機(jī)變量的乘積的期望等于它們的期望的乘積。
$D(X)=E(X^{2}) - [E(X)]^{2}$：一個(gè)隨機(jī)變量的方差等于它的平方的期望減去期望的平方
$D(X+Y)=D(X)+D(Y)$：兩個(gè)相互獨(dú)立的隨機(jī)變量之和的方差等于它們的方差的和。

可以推導(dǎo)出兩個(gè)隨機(jī)變量的乘積的方差如下：

$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}$

又由于輸入變量是符合標(biāo)準(zhǔn)的正態(tài)分布的，因此 $E(X)=0$，$E(Y)=0$，可知 $D(X \times Y)=D(X) \times D(Y)$

我們以輸入層第一個(gè)神經(jīng)元為例：

$\mathrm{H}_{11}=\sum_{i=0}^{n} X_{i} * W_{1 i}$

其中輸入 X 和權(quán)值 W 都是服從 $N(0,1)$ 的正態(tài)分布，且由公式$D(X \times Y)=D(X) \times D(Y)$, 因此這個(gè)神經(jīng)元的方差為：

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

可以求其標(biāo)準(zhǔn)差：

$\operatorname{std}\left(\mathrm{H}_{11}\right)=\sqrt{\mathrm{D}\left(\mathrm{H}_{11}\right)}=\sqrt{n}$

可見，經(jīng)過第一層網(wǎng)絡(luò)，方差就會(huì)擴(kuò)大 n 倍，標(biāo)準(zhǔn)差就擴(kuò)大 $\sqrt{n}$ 倍，n 為每層神經(jīng)元個(gè)數(shù)，直到超出數(shù)值表示范圍。

從前面的輸出中也可以看出來，n = 256，因此每一層的標(biāo)準(zhǔn)差輸出都是16倍。再由公式可知，每一層網(wǎng)絡(luò)輸出的方差與神經(jīng)元個(gè)數(shù)、輸入數(shù)據(jù)的方差、權(quán)值方差有關(guān)（見上式），通過觀察可知，比較好改變的是權(quán)值的方差 $D(W)$，要控制每一層輸出的方差仍然為1左右，因此需要 $D(W)= \frac{1}{n}$，可知標(biāo)準(zhǔn)差為 $std(W)=\sqrt\frac{1}{n}$。因此修改權(quán)值初始化代碼為nn.init.normal_(m.weight.data, std=np.sqrt(1/self.neural_num))

再次輸出時(shí)，結(jié)果如下：

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

修改之后，沒有出現(xiàn)梯度消失或者梯度爆炸的情況，每層神經(jīng)元輸出的方差均在 1 左右。通過恰當(dāng)?shù)臋?quán)值初始化，可以保持權(quán)值在更新過程中維持在一定范圍之內(nèi)。

但是上述的實(shí)驗(yàn)前提為未使用非線性函數(shù)的前提下，如果在forward()中添加非線性變換例如tanh，每一層的輸出方差會(huì)越來越小，會(huì)導(dǎo)致梯度消失。

為了解決這個(gè)問題，進(jìn)一步有了著名的 Xavier 初始化與 Kaiming 初始化。

Xavier 方法與 Kaiming 方法

Xavier 方法

Xavier 是 2010 年提出的，針對(duì)有非線性激活函數(shù)時(shí)的權(quán)值初始化方法。

目標(biāo)是保持?jǐn)?shù)據(jù)的方差維持在 1 左右
針對(duì)飽和激活函數(shù)如 sigmoid 和 tanh 等。

同時(shí)考慮前向傳播和反向傳播，需要滿足兩個(gè)等式

$\begin{array}{l}
\boldsymbol{n}_{\boldsymbol{i}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1} \
\boldsymbol{n}_{\boldsymbol{i}+\mathbf{1}} * \boldsymbol{D}(\boldsymbol{W})=\mathbf{1} \
\end{array}$

通過計(jì)算可知：

$D(W)=\frac{2}{n_{i}+n_{i+1}}$。

為了使 Xavier 方法初始化的權(quán)值服從均勻分布，假設(shè) $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，計(jì)算增益
tanh_gain = nn.init.calculate_gain('tanh')
a *= tanh_gain
nn.init.uniform_(m.weight.data, -a, a)

并且每一層的激活函數(shù)都使用 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 左右。

也可以直接調(diào)用PyTorch 中 Xavier 初始化方法：

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

nn.init.calculate\_gain()

這里重點(diǎn)介紹一下nn.init.calculate_gain(nonlinearity,param=**None**)方法。

主要功能是經(jīng)過一個(gè)分布的方差經(jīng)過激活函數(shù)后的變化尺度，主要有兩個(gè)參數(shù)：

nonlinearity：激活函數(shù)名稱
param：激活函數(shù)的參數(shù)，如 Leaky ReLU 的 negative\_slop等等。

下面是計(jì)算標(biāo)準(zhǔn)差經(jīng)過激活函數(shù)的變化尺度的代碼。

x = torch.randn(10000)
out = torch.tanh(x)
# 計(jì)算變化尺度（也可以稱為變化倍數(shù)）
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

結(jié)果表示，原有數(shù)據(jù)分布的方差經(jīng)過 tanh 之后，標(biāo)準(zhǔn)差會(huì)變小 1.6 倍左右。

Kaiming 方法

雖然 Xavier 方法提出了針對(duì)飽和激活函數(shù)的權(quán)值初始化方法，但是 AlexNet 出現(xiàn)后，大量網(wǎng)絡(luò)開始使用非飽和的激活函數(shù)如 ReLU 等，這時(shí) Xavier 方法不再適用。2015 年針對(duì) ReLU 及其變種等激活函數(shù)提出了 Kaiming 初始化方法。

針對(duì) ReLU，方差應(yīng)該滿足：$\mathrm{D}(W)=\frac{2}{n_{i}}$；

針對(duì) ReLu 的變種，方差應(yīng)該滿足：$D(W)=\frac{2}{\left(1+\mathrm{a}^{2}\right) * n_{i}}$，a 表示負(fù)半軸的斜率，如 PReLU 方法，標(biāo)準(zhǔn)差滿足 $\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)。