【深度学习】PyTorch框架(3)：优化与初始化

1.引言

在本文中，我们将探讨神经网络的优化与初始化技术。随着神经网络深度的增加，我们会遇到多种挑战。最关键的是确保网络中梯度流动的稳定性，否则可能会遭遇梯度消失或梯度爆炸的问题。因此，我们将深入探讨以下两个核心概念：网络参数的初始化和优化算法的选择。
本文的前半部分，我们将介绍不同的参数初始化方法，从最基本的初始化策略开始，逐步深入到当前在极深网络中应用的高级技术。在后半部分，我们将聚焦于优化算法的比较，分析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') # 用于导出
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
#我们将使用与教程3相同的set_seed函数，以及路径变量DATASET_PATH和CHECKPOINT_PATH。如有必要，请调整路径。

# 数据集下载存放的文件夹路径（例如MNIST）
DATASET_PATH = "../data"
# 预训练模型保存的文件夹路径
CHECKPOINT_PATH = "../saved_models/tutorial4"

# 设置种子的函数
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

# 获取将在此笔记本中使用整个过程中使用的设备
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
print("Using device", device)
使用设备 cuda:0
##在本文的最后部分，我们将使用三种不同的优化器训练模型。以下是这些模型的预训练版本下载链接。

import urllib.request
from urllib.error import HTTPError
# 存储本教程预训练模型的Github URL
base_url = "https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial4/" 
# 需要下载的文件
pretrained_files = ["FashionMNIST_SGD.config",    "FashionMNIST_SGD_results.json",    "FashionMNIST_SGD.tar",
                    "FashionMNIST_SGDMom.config", "FashionMNIST_SGDMom_results.json", "FashionMNIST_SGDMom.tar",
                    "FashionMNIST_Adam.config",   "FashionMNIST_Adam_results.json",   "FashionMNIST_Adam.tar"   ]
# 如果检查点路径不存在，则创建
os.makedirs(CHECKPOINT_PATH, exist_ok=True)

# 对于每个文件，检查它是否已经存在。如果不存在，尝试下载。
for file_name in pretrained_files:
    file_path = os.path.join(CHECKPOINT_PATH, file_name)
    if not os.path.isfile(file_path):
        file_url = base_url + file_name
        print(f"正在下载 {
     
     file_url}...")
        try:
            urllib.request.urlretrieve(file_url, file_path)
        except HTTPError as e:
            print("下载过程中出现问题。请尝试从GDrive文件夹下载文件，或联系作者，并附上包括以下错误的完整输出：\n", e)

2.准备工作

在本文中，我们将使用一个深度全连接网络，与我们之前的文章类似。我们还将再次将网络应用于FashionMNIST，我们首先加载FashionMNIST数据集：

from torchvision.datasets import FashionMNIST
from torchvision import transforms

# 应用于每张图片的转换 => 首先将它们转换为张量，然后使用均值为0和标准差为1进行归一化
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.2861,), (0.3530,))
])

# 加载训练数据集。我们需要将其分割为训练部分和验证部分
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)

与之前的文章相比，我们更改了归一化转换transforms.Normalize的参数。现在归一化的设计是让我们在像素上获得预期的均值为0和标准差为1。这将特别适用于我们下面将要讨论的初始化问题，因此我们在这里进行更改。应当指出，在大多数分类任务中，两种归一化技术（介于-1和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包含许多黑色像素，与MNIST类似。接下来，我们将创建一个线性神经网络。

class BaseNetwork(nn.Module):
    def __init__(self, act_fn, input_size=784, num_classes=10, hidden_sizes=[512, 256, 256, 128]):
        """
        输入:
            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) # 模块列表将模块列表注册为子模块（例如，用于参数）

        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
}

最后，我们定义了一些绘图函数，我们将在讨论中使用它们。这些函数帮助我们
(1)可视化网络内部的权重/参数分布，
(2)可视化不同层的参数接收的梯度，以及
(3)激活值，即线性层的输出。

# 绘制值的分布图
def plot_dists(val_dict, color="C0", xlabel=None, stat="count", use_kde=True):
    columns = len(val_dict)  # 图表的列数等于val_dict的键的数量
    fig, ax = plt.subplots(1, columns, figsize=(columns*3, 2.5))  # 创建子图
    fig_index = 0
    for key in sorted(val_dict.keys()):  # 遍历val_dict的键
        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))  # 如果有方差则绘制核密度估计
        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)  # 设置x轴标签
        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()  # 将权重转换为numpy数组

    # 绘图
    fig = plot_dists(weights, color=color, xlabel="Weight vals")  # 使用plot_dists函数绘制权重分布图
    fig.suptitle("Weight distribution", fontsize=14, y=1.05)  # 设置图表标题
    plt.show()  # 显示图表
    plt.close()  # 关闭图表

# 可视化模型梯度分布
def visualize_gradients(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)  # 将数据移动到设备上

    # 将一批数据通过网络前向传播，并计算权重的梯度
    model.zero_grad()  # 清空梯度
    preds = model(imgs)  # 前向传播
    loss = F.cross_entropy(preds, labels)  # 计算交叉熵损失
    loss.backward()  # 反向传播计算梯度
    # 限制可视化为权重参数，不包括偏置，以减少图表数量
    grads = {
   
   name: params.grad.view(-1).cpu().clone().numpy() for name, params in model.named_parameters() if "weight" in name}
    model.zero_grad()  # 清空梯度

    # 绘图
    fig = plot_dists(grads, color=color, xlabel="Grad magnitude")  # 使用plot_dists函数绘制梯度分布图
    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)  # 将数据移动到设备上

    # 将一批数据通过网络前向传播，并计算权重的梯度
    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()  # 将激活值转换为numpy数组

    # 绘图
    fig = plot_dists(activations, color=color, stat="density", xlabel="Activation vals")  # 使用plot_dists函数绘制激活值分布图
    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])}")  # 打印方差

3.初始化

在深入讨论神经网络的初始化问题之前，有必要指出，关于这一主题，网络上已经有许多精彩的博客文章，例如deeplearning.ai提供的资源，或者那些更侧重于数学分析的文章。如果在阅读完本教程后仍有疑惑，我们建议您也浏览一下这些博客文章以获得更深入的理解。
初始化神经网络时，我们希望其具备一些特定的属性。首先，输入数据的方差应能通过整个网络传递到输出层，以保证输出神经元具有相似的标准差。如果我们在网络深层发现方差逐渐消失，那么模型将难以优化，因为下一层的输入将变得几乎等同于一个恒定值。同样地，如果方差随着网络深度的增加而增大，那么梯度可能会变得非常大，导致数值稳定性问题。其次，我们希望在初始化时各层的梯度分布具有相同的方差。如果第一层得到的梯度远小于最后一层，我们可能就会在选择合适的学习速率时遇到困难。
为了寻找合适的初始化方法，我们首先以一个没有激活函数的线性神经网络作为起点进行分析，即网络中仅使用恒等激活函数。之所以这样做，是因为不同的激活函数对初始化方法有特定的要求，我们可以根据所使用的激活函数调整初始化策略。

model = BaseNetwork(act_fn=Identity()).to(device)

3.1 常数初始化

接下来，我们考虑一种最简单的初始化方法——常数初始化。直观上，将所有权重设置为零并不理想，因为这会导致传播的梯度也为零。但是，如果我们将所有权重设置为一个接近零的非零常数，情况会如何呢？为了探究这一点，我们可以编写一个函数来实现这一初始化，并可视化梯度的分布情况。
定义了一个名为const_init的函数，它接受一个模型和一个默认为0的常数值c，将模型中所有的参数（权重）填充为这个常数值。然后，我们调用这个函数将模型的权重初始化为0.005，接着使用visualize_gradients和visualize_activations函数来可视化梯度和激活值的分布，并打印出它们的方差。这有助于我们理解在这种初始化策略下，网络的梯度和激活值的行为。

def const_init(model, c=0.0):
    for name, param in model.named_parameters():
        param.data.fill_(c)

const_init(model, c=0.005)
visualize_gradients(model)
visualize_activations(model, print_variance=True)

Layer 0 - Variance: 2.058276
Layer 2 - Variance: 13.489119
Layer 4 - Variance: 22.100567
Layer 6 - Variance: 36.209572
Layer 8 - Variance: 14.831439

从我们的观察来看，只有第一层和最后一层展现出了多样化的梯度分布，而中间的三层则显示出所有权重具有相同梯度的现象（注意，这个值并不为零，但往往非常接近零）。如果用相同值初始化的参数最终获得了相同的梯度，这就意味着这些参数的值将始终一致。这样的结果会让我们网络中的这一层失去作用，实际上将我们网络的参数数量减少到了单一的一个值。因此，我们不能采用常数值初始化的方法来训练我们的网络。

3.2.关于方差的恒定性

在上述实验中，我们已经发现单一的常数值初始化策略是行不通的。那么，如果我们改为从诸如高斯分布的某种概率分布中随机采样来初始化参数，情况会怎样呢？最直接的方法可能是为网络中的所有层选择一个相同的方差值。接下来，我们将实现这种方法，并可视化各层的激活分布情况。

def var_init(model, std=0.01):
    for name, param in model.named_parameters():
        param.data.normal_(std=std)

var_init(model, std=0.01)
visualize_activations(model, print_variance=True)

在神经网络的层与层之间，激活值的方差呈现出逐渐减小的趋势，到了最后一层，方差几乎趋近于零。这种情况下，一个可能的解决办法是增加标准差的数值。通过提高初始化时的标准差，我们可以尝试维持网络深层的激活方差，避免其在传播过程中消失。

var_init(model, std=0.1)
visualize_activations(model, print_variance=True)

通过使用更高的标准差进行初始化，我们可以观察到网络各层激活值的分布情况，特别是它们的方差，以评估这种策略是否有效。这种方法可能有助于解决深层网络中的梯度消失问题，但同时也要警惕不要导致梯度爆炸，这需要我们在实践中仔细调整和平衡。

3.3.如何找到合适的初始化值

从我们之前的实验中，我们可以看到需要从某个概率分布中对权重进行采样，但具体选择哪个分布我们还不确定。下一步，我们将尝试从激活值分布的角度出发，寻找最优的初始化方法。为此，我们提出两个要求：
（1）激活值的均值应该为零。
（2）激活值的方差应该在每一层都保持不变。
假设我们要为以下层设计一个初始化方法：要设计一个满足上述两个要求的初始化方法，我们需要考虑权重矩阵和激活函数的特性。对于一个全连接层，如果使用恒等激活函数（即线性激活函数），输出的均值和方差将取决于输入的均值和方差以及权重矩阵。为了使激活值的均值为零，我们可以选择一个合适的均值。而为了保持激活值的方差在每一层都相同，我们需要选择一个合适的标准差。
一种常见的方法是使用与输入维度的平方根成反比的标准差。这样，无论输入的维度如何变化，权重的标准差都会相应调整，以保持激活值的方差大致相同。这种方法通常被称为Xavier初始化或Glorot初始化。
我们的目标是让每个元素的方差与输入的方差相同，即每个权重更新后的方差应该保持与输入数据的方差一致。这是为了确保在多层网络中，信息能够稳定地从前层传递到后层，避免出现梯度消失或爆炸的问题。

在数学上，如果我们考虑一个全连接层$y=Wx+b,y\in {\mathbb{R}}^{d_y},x\in {\mathbb{R}}^{d_{}}$