深入浅出 Pytorch 系列 — 参数初始化(1)携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的

携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第14天，点击查看活动详情

文章可能写有些仓促，细节可能不算到位，先发表，随后会不定期更新和补充内容

今天我们主要来从新审视一下神经网络的参数优化和初始化。随着神经网络层数不断增加，我们就会遇到各种各样因为层数增加所带来的问题。最重要就是我们需要一个稳定的梯度流，来更新参数，如果发生了梯度消失和梯度爆炸。为了预防这些问题发生，我们就有必要来看一看参数初始化和优化。

首先我们来尝试各种参数初始化的方法，我们先从简单开始，然后一步一步加大难度，通过结果来分析每种初始化方式对训练影响。随后会还会讨论优化器对训练过程的影响，通过比较 SGD、带动量的 SGD 以及 Adam 来查看这些优化器对训练带来哪些影响。

引入依赖模块

标准库

import os
import json
import math
import numpy as np
import copy

绘制相关库

import matplotlib.pyplot as plt
from matplotlib import cm
%matplotlib inline
from IPython.display import set_matplotlib_formats
set_matplotlib_formats('svg', 'pdf') # For export
import seaborn as sns
sns.set()

进度条

from tqdm.notebook import tqdm

PyTorch

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim

# 保存数据集的目录(e.g. MNIST)
DATASET_PATH = "/data"
# 保存预训练模型的目录
CHECKPOINT_PATH = "../saved_models/tutorial4"

# Function for setting the seed
def set_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
set_seed(42)

用于确保所有在 GPU 上运算都是一致性

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

遍历可用的设备，如果有可用 GPU 首选在 GPU 设备上运行代码

device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
print("Using device", device)

准备数据集

和之前关于激活函数分享类似，这里也使用一个全连接网络，数据集依旧采用的是 FashionMNIST。

from torchvision.datasets import FashionMNIST
from torchvision import transforms

变换

对数据集应用变换，首先将数据集数据转换为 tensor,然后将数据集图像像素值进行归一化为标准正态分布，均值为 0 方差为 1。

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.2861,), (0.3530,))
                               ])

下载训练数据集，然后将训练数据集拆分为训练集和验证集 2 个部分

train_dataset = FashionMNIST(root=DATASET_PATH, train=True, transform=transform, download=True)
train_set, val_set = torch.utils.data.random_split(train_dataset, [50000, 10000])

下载测试数据集

test_set = FashionMNIST(root=DATASET_PATH, train=False, transform=transform, download=True)

我们定义一系列数据加载器，这些数据加载用于加载训练数据、验证数据或者测试数据。值得注意的是在实际训练过程会采用小批量的


train_loader = data.DataLoader(train_set, batch_size=1024, shuffle=True, drop_last=False)
val_loader = data.DataLoader(val_set, batch_size=1024, shuffle=False, drop_last=False)
test_loader = data.DataLoader(test_set, batch_size=1024, shuffle=False, drop_last=False)

与之前的分享相比，这里数据集标准化定义有所不同，并不是简单地给出(0.5,0.5) 我们期望通过标准化后，图像像素分布是是均值为 0 标准差为 1 的这样的分布。

标准化是将图像像素值标准化到均值为 0 标准差为 1 的范围，这个与随后我们讨论的初始有关，通过计算图像像素的均值和标准，然后用均值和标准差对数据集图像做标准化。

print("Mean", (train_dataset.data.float() / 255.0).mean().item())
print("Std", (train_dataset.data.float() / 255.0).std().item())

Mean 0.2860923707485199
Std 0.3530242443084717

imgs, _ = next(iter(train_loader))
print(f"Mean: {imgs.mean().item():5.3f}")
print(f"Standard deviation: {imgs.std().item():5.3f}")
print(f"Maximum: {imgs.max().item():5.3f}")
print(f"Minimum: {imgs.min().item():5.3f}")

Mean: 0.002
Standard deviation: 1.001
Maximum: 2.022
Minimum: -0.810

注意最大值和最小值并不是 1 和 -1 ，这个可能是因为 FashionMNIST 中包含一些 black 像素值，整个分布向整数方向有所偏移

class BaseNetwork(nn.Module):

    def __init__(self, act_fn, input_size=784, num_classes=10, hidden_sizes=[512, 256, 256, 128]):
        """
        Inputs:
            act_fn - 激活函数的对象，作为神经网络非线性部分
            input_size - 输入图像尺寸(单位为像素)
            num_classes - 要预测的类别数
            hidden_sizes - 隐藏层神经元数量列表
        """
        super().__init__()

        layers = []
        layer_sizes = [input_size] + hidden_sizes
        for layer_index in range(1, len(layer_sizes)):
            layers += [nn.Linear(layer_sizes[layer_index-1], layer_sizes[layer_index]),
                       act_fn]
        layers += [nn.Linear(layer_sizes[-1], num_classes)]
        self.layers = nn.ModuleList(layers) # A module list registers a list of modules as submodules (e.g. for parameters)

        self.config = {"act_fn": act_fn.__class__.__name__, "input_size": input_size, "num_classes": num_classes, "hidden_sizes": hidden_sizes}

    def forward(self, x):
        x = x.view(x.size(0), -1)
        for l in self.layers:
            x = l(x)
        return x

对于激活函数，这一次使用 Pytorch 提供的 torch.nn 包中提供激活函数，而不再是自己去实现激活函数。不过这里定义了一个 Identity 激活函数，虽然这个激活函数在很大程度上限制模型的能力，使用这个激活目的在于将问题简化便于研究初始化。

class Identity(nn.Module):
    def forward(self, x):
        return x

act_fn_by_name = {
    "tanh": nn.Tanh,
    "relu": nn.ReLU,
    "identity": Identity
}

最后，我们来定义几个绘制工具，用于将一些数据以图表形式展示出来

查看权重/参数的分布
查看各个层参数的梯度
激活层输出

关于这些绘制工具具体实现在下面一一列出，大家如果感兴趣可以看一下

绘制分布

def plot_dists(val_dict, color="C0", xlabel=None, stat="count", use_kde=True):
    columns = len(val_dict)
    fig, ax = plt.subplots(1, columns, figsize=(columns*3, 2.5))
    fig_index = 0
    for key in sorted(val_dict.keys()):
        key_ax = ax[fig_index%columns]
        sns.histplot(val_dict[key], ax=key_ax, color=color, bins=50, stat=stat,
                     kde=use_kde and ((val_dict[key].max()-val_dict[key].min())>1e-8)) # Only plot kde if there is variance
        key_ax.set_title(f"{key} " + (r"(%i $\to$ %i)" % (val_dict[key].shape[1], val_dict[key].shape[0]) if len(val_dict[key].shape)>1 else ""))
        if xlabel is not None:
            key_ax.set_xlabel(xlabel)
        fig_index += 1
    fig.subplots_adjust(wspace=0.4)
    return fig

绘制权重分布

def visualize_weight_distribution(model, color="C0"):
    weights = {}
    for name, param in model.named_parameters():
        if name.endswith(".bias"):
            continue
        key_name = f"Layer {name.split('.')[1]}"
        weights[key_name] = param.detach().view(-1).cpu().numpy()

    ## Plotting
    fig = plot_dists(weights, color=color, xlabel="Weight vals")
    fig.suptitle("Weight distribution", fontsize=14, y=1.05)
    plt.show()
    plt.close()

绘制梯度

def visualize_gradients(model, color="C0", print_variance=False):
    """
    Inputs:
        net - Object of class BaseNetwork
        color - Color in which we want to visualize the histogram (for easier separation of activation functions)
    """
    model.eval()
    small_loader = data.DataLoader(train_set, batch_size=1024, shuffle=False)
    imgs, labels = next(iter(small_loader))
    imgs, labels = imgs.to(device), labels.to(device)

    # Pass one batch through the network, and calculate the gradients for the weights
    model.zero_grad()
    preds = model(imgs)
    loss = F.cross_entropy(preds, labels) # Same as nn.CrossEntropyLoss, but as a function instead of module
    loss.backward()
    # We limit our visualization to the weight parameters and exclude the bias to reduce the number of plots
    grads = {name: params.grad.view(-1).cpu().clone().numpy() for name, params in model.named_parameters() if "weight" in name}
    model.zero_grad()

    ## Plotting
    fig = plot_dists(grads, color=color, xlabel="Grad magnitude")
    fig.suptitle("Gradient distribution", fontsize=14, y=1.05)
    plt.show()
    plt.close()

    if print_variance:
        for key in sorted(grads.keys()):
            print(f"{key} - Variance: {np.var(grads[key])}")

绘制激活函数

def visualize_activations(model, color="C0", print_variance=False):
    model.eval()
    small_loader = data.DataLoader(train_set, batch_size=1024, shuffle=False)
    imgs, labels = next(iter(small_loader))
    imgs, labels = imgs.to(device), labels.to(device)

    # Pass one batch through the network, and calculate the gradients for the weights
    feats = imgs.view(imgs.shape[0], -1)
    activations = {}
    with torch.no_grad():
        for layer_index, layer in enumerate(model.layers):
            feats = layer(feats)
            if isinstance(layer, nn.Linear):
                activations[f"Layer {layer_index}"] = feats.view(-1).detach().cpu().numpy()

    ## Plotting
    fig = plot_dists(activations, color=color, stat="density", xlabel="Activation vals")
    fig.suptitle("Activation distribution", fontsize=14, y=1.05)
    plt.show()
    plt.close()

    if print_variance:
        for key in sorted(activations.keys()):
            print(f"{key} - Variance: {np.var(activations[key])}")

首先我们需要简化问题，为了更好分析线性神经网络初始化，我们先采用 Identity 这个激活函数，好处是让神经网络不会受到激活函数干扰，因不同激活函数对于初始化可能存在不一样反应，也就是不同激活函数和参数初始化方式不同会产生不同效果。随后可以根据激活函数选择来采用不同初始化。

屏幕快照 2022-08-07 下午1.11.42.png

对于深度学习框架，不同于我们熟知开发软件，更多时候，很少会有逻辑这样 bug，未达到预期效果背后原因有很多、网络结构、数据集、参数的初始化、优化器的选择和目标函数定义等等。

参数初始化

早期参数初始化都是将数据和参数做均值为 0 方差为 1 标准化处理。随着神经网络层数加深，这样方式来初始化参数往往并不能解决梯度消失和梯度爆炸的问题。

激活值的方差是逐层递减的，这导致反向传播中的梯度也逐层递减。要解决梯度消失，就要避免激活值方差的衰减，最理想的情况是，每层的输出值（激活值）保持正态分布。

整个大型前馈神经网络无非就是一个超级大映射，将原始样本稳定的映射成其的类别。也就是将样本空间映射到类别空间。试想，如果样本空间与类别空间的分布差异很大，

类别空间特别稠密，样本空间特别稀疏辽阔，那么在类别空间得到的用于反向传播的误差丢给样本空间后简直变得微不足道，也就是会导致模型的训练非常缓慢。
类别空间特别稀疏，样本空间特别稠密，那么在类别空间算出来的误差丢给样本空间后简直是爆炸般的存在，即导致模型发散震荡，无法收敛。因此，我们要让样本空间与类别空间的分布差异（密度差别）不要太大，也就是要让它们的方差尽可能相等。

对方差的性质进行简单回顾

在开始之前，我们先简单复习关于方差和均值的公式，当 $X$ 和 $Y$ 相互独立

Var(X + Y) = Var(X) + Var(Y)

方差的计算公式

Var(X) = \mathbb{E}(X^2) - \mathbb{E}(X)^2

为了说明问题，我们这里神经网络只有一层，且该层只有一个神经元，并且暂时不考虑偏置所以就有

y=w_1x_1 + w_2x_2 + w_3x_3

在训练神经网络之前，我们需要对参数进行初始化，那么应该如何选择参数，才能让神经网络顺利训练起来。

屏幕快照 2022-08-08 下午8.10.52.png

关于网络参数初始化，我们可以将 $w = [0,0,0]$ 也就是将 w 初始化为全 0 的向量。那么如果我们原来基础上给这个层再去添加一个神经元，那么现在就有了两个神经元，并且假设这两个神经元权重都是 0 那么也就是

w_1 = [w_{11},w_{12},w_{13}] = [0,0,0]\\ w_2 = [w_{21},w_{22},w_{23}] = [0,0,0]\\

屏幕快照 2022-08-08 下午8.15.48.png

然后我们去计算梯度 $w_{11}$ 和 $w_{21}$

\frac{\partial L}{\partial w_{11}} = \frac{\partial L}{\partial y_1} x_1

\frac{\partial L}{\partial w_{12}} = \frac{\partial L}{\partial y_2} x_1

不难看出因为参数一样以及输入值一样所以 $y_1$ 和 $y_2$ 有一样值，计算梯度也是具有相同值，那么这两个神经元具有相同权重初始值，在反向传播时，梯度也是一样的，那么这两个神经元的变化在学习中也是始终保持一致，这时神经元就是失去了差异性，从而无法通过这两神经元来学习到表达更复杂的特征

一般我们都会用一个正态分布来随机初始化参数，那么为了简化问题，我们将输入 $x_1,x_2,x_3$ 假设为 1 那么就得到了

y = w_1 + w_2 + w_3

因为 $w_1,w_2,w_3$ 是独立同分布的采样，根据上面独立随机变量方差求和公式可以得到下面式子

Var(y) = Var(w_1) + Var(w_2) + Var(w_3)

从而可以计算 y 的标准差为

std(y) = \sqrt{3} = 1.732

这也就是表示输入经过神经元的激活函数之后，对输入离散程度被提高了，如果现在神经元是 n 输入那么经过激活函数输出的方差就变为原有 n 倍，我们希望输入方差和输出方差是相等，依次作为目标来来设置参数。

y = \sum_{i=1}^n w_i\\ Var(y) = n Var(w_i) = n\\ Std(y) = \sqrt{n}

Xavier 参数初始化

本着让输入方差和输出方差相等原则来设计参数 w 的方差和均值，对于参数 w 均值设置为 0 那么主要是看其方差应该如何选择。为了让输入方差和输出方差都为 1。

Var(y) = nVar(w_i) = 1

这样求出 $Var(w_i) = \frac{1}{n}$ 如果我们考虑输入层，那么 $w_i$ 方差就等于 1/n

Var(w_i) = \frac{2}{n_{in} + n_{out}}

正态分布初始化

\cal{N}(0,\frac{2}{n_{in} + n_{out}})

这个正态分布中随机采样即可

W1 = torch.randn(784, nh) * math.sqrt(1 / 784)
b1 = torch.zeros(nh)
W2 = torch.randn(nh, 1) * math.sqrt(1 / nh)
b2 = torch.zeros(1)

z1 = linear(x_train, W1, b1)
print(z1.mean(), z1.std())

tensor(0.1031) tensor(0.9458)

a1 = relu(z1)
a1.mean(), a1.std()

(tensor(0.4272), tensor(0.5915))

均匀分布

当 x 满足从 a 到 b 的均匀分布时，那么 x 方差等于

Var(x) = \mathbb{E}(x^2) - \mathbb{E}(x)^2 = \frac{(b - a)^2}{12}

为了保证采样均值为 0 那么可以将采用写成从 -a 到 a 的均匀分布，那么现在方差也就是

Var(x) = \frac{a^2}{3}

那么现在如果我们将目标方差带入公式

Var(w_i) = \frac{2}{n_{in} + n_{out}}

\cal{u}\left(- \sqrt{\frac{6}{n_{in} + n_{out}}} ,\sqrt{\frac{6}{n_{in} + n_{out}}} \right)

Xavier 初始化（亦称Glorot 初始化）最初由Xavier Glorot 等人在2010 的 “Understanding the difficulty of training deep feedforward neural networks” 一文中提出，其核心思想是使层的输出数据的方差与其输入数据的方差相等。而在 Xavier 初始化的论文中，作者只考虑了当时默认的 logistic sigmoid 激活函数。

kaiming 参数初始化

因为relu会抛弃掉小于0的值，对于一个均值为0的data来说，这就相当于砍掉了一半的值，这样一来，均值就会变大，前面Xavier初始化公式中E(x)=mean=0的情况就不成立了。根据新公式的推导，最终得到新的rescale系数

因为经过 ReLU 激活函数，大概有一半输出会变成 0

Var(y) = \frac{1}{2}nVar(w_i) = 1

那么我们目标方差就变为

Var(w_i) = \frac{2}{n}

kaiming 参数初始化

如果我们用的是 LeakyReLU 这样激活函数那么

Var(w_i) = \frac{2}{n(1 + \alpha^2)}

正态初始化

\cal{N}(0,\frac{2}{n(1 + \alpha^2)})

均匀分布初始化

\cal{u}\left(- \sqrt{\frac{6}{n(1 + \alpha^2)}} ,\sqrt{\frac{6}{n(1 + \alpha^2)}} \right)

W1 = torch.randn(784, nh) * math.sqrt(2 / 784)
b1 = torch.zeros(nh)
W2 = torch.randn(nh, 1) * math.sqrt(2 / nh)
b2 = torch.zeros(1)

z1 = linear(x_train, W1, b1)
a1 = relu(z1)
a1.mean(), a1.std()

(tensor(0.4553), tensor(0.7339))