PyTorch中的权重初始化方法

恰当的权重初始化是搭建神经网络以实现有效训练的重要一步。模型权重的初始值会极大影响其学习能力。选择不当的初始权重可能导致梯度消失或爆炸、收敛缓慢，或者网络陷入次优局部极小值。如果您使用过 Keras，可能熟悉为层指定 glorot_uniform 或 he_normal 等初始化器。PyTorch 提供了类似的灵活性，不过它应用这些初始化方式的方法更明确。

本节将带领您了解 PyTorch 中常见的权重初始化方法，以及如何将它们应用于基于 torch.nn.Module 的模型。我们将与 TensorFlow Keras 的实践进行比较，以帮助您迁移现有知识。

良好初始化的作用

在网络开始从数据中学习之前，其权重需要设置为一些初始值。请思考以下几点:

• 打破对称性: 如果所有权重都初始化为相同的值（例如，都为零），那么层中的所有神经元在反向传播过程中将产生相同的输出并计算相同的梯度。这种对称性会阻止网络学习到不同的特征。随机初始化有助于打破这种对称性。
• 避免梯度问题: 训练期间，梯度通过网络反向传播。如果权重过小，梯度会呈指数级减小（梯度消失），使早期层难以学习。相反，如果权重过大，梯度会呈指数级增长（梯度爆炸），导致训练不稳定。
• 加快收敛速度: 周全的初始化方法可以将权重放置在优化更有效的区域，从而加快收敛速度。

PyTorch 层中的默认初始化

当您在 PyTorch 中创建一个层时，例如 nn.Linear 或 nn.Conv2d，其参数（权重和偏置）会自动初始化。例如，nn.Linear 和 nn.Conv2d 层默认对其权重使用 Kaiming 均匀初始化。偏置（如果存在）通常初始化为零，或者使用源自 Kaiming 初始化范围的均匀分布。

虽然这些默认设置通常是合理的起点，特别是对于常见架构而言，但您通常会想应用针对您的网络架构或激活函数定制的特定初始化方案。

torch.nn.init 模块

PyTorch 将其初始化函数集中在 torch.nn.init 模块中。此模块中的大多数函数都是就地操作，这意味着它们直接修改输入张量。这是 PyTorch 的常见约定，通常由函数名中的下划线后缀表示（例如，xavier_uniform_）。

我们来看看一些最常用的初始化器。

均匀分布和正态分布

最简单的初始化器从标准分布中抽取值：

• nn.init.uniform_(tensor, a=0.0, b=1.0): 用从均匀分布 $U(a, b)$ 中抽取的值填充输入 tensor。
• nn.init.normal_(tensor, mean=0.0, std=1.0): 用从正态分布 $N(\text{均值}, \text{标准差}^2)$ 中抽取的值填充输入 tensor。

在 Keras 中，您可以使用 tf.keras.initializers.RandomUniform 或 tf.keras.initializers.RandomNormal。

常量初始化（零、一、自定义值）

有时，您需要将权重或偏置初始化为特定的常量值：

• nn.init.zeros_(tensor): 用零填充输入 tensor。通常用于初始化偏置。
• nn.init.ones_(tensor): 用一填充输入 tensor。
• nn.init.constant_(tensor, val): 用特定值 val 填充输入 tensor。

Keras 的对应项包括 tf.keras.initializers.Zeros、tf.keras.initializers.Ones 和 tf.keras.initializers.Constant。

Xavier (Glorot) 初始化

Xavier 初始化由 Glorot 和 Bengio（2010 年）提出，旨在使激活值和梯度的方差在各层间大致保持不变。这有助于防止信号衰减或爆炸。它特别适合于后面接有 sigmoid 或 tanh 等激活函数的层。

• nn.init.xavier_uniform_(tensor, gain=1.0): 根据均匀分布填充输入 tensor 的值，其范围基于 tensor 的 fan_in（输入单元的数量）和 fan_out（输出单元的数量）以及可选的 gain 计算得出。范围是 $[-界限, 界限]$，其中 $$\text{界限} = \text{增益} \times \sqrt{\frac{6}{fan\_in + fan\_out}}$$
• nn.init.xavier_normal_(tensor, gain=1.0): 用从正态分布中抽取的值填充输入 tensor，其标准差为 $$\text{标准差} = \text{增益} \times \sqrt{\frac{2}{fan\_in + fan\_out}}$$

gain 参数允许调整以适应不同的激活函数；nn.init.calculate_gain(nonlinearity, param=None) 可用于确定合适的增益（例如，nn.init.calculate_gain('relu')）。 在 Keras 中，这些是 tf.keras.initializers.GlorotUniform 和 tf.keras.initializers.GlorotNormal。

Kaiming (He) 初始化

Kaiming 初始化由 He 等人（2015 年）提出，专门为后面接有修正线性单元 (ReLU) 激活函数及其变体（如 Leaky ReLU）的层设计。它考虑了 ReLU 将负输入设置为零这一事实，这会影响输出的方差。

• nn.init.kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'): 使用均匀分布。模式 可以是 'fan_in'（在正向传播中保持方差）或 'fan_out'（在反向传播中保持方差）。非线性 参数指定层后使用的激活函数。 均匀分布的界限是 $[-界限, 界限]$，其中 $$\text{界限} = \sqrt{\frac{6}{(1 + a^2) \times fan\_mode}}$$ 这里，$a$ 是 Leaky ReLU 的负斜率（标准 ReLU 为 0）。
• nn.init.kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'): 使用正态分布，其标准差为 $$\text{标准差} = \sqrt{\frac{2}{(1 + a^2) \times fan\_mode}}$$

Keras 提供 tf.keras.initializers.HeUniform 和 tf.keras.initializers.HeNormal。如前所述，PyTorch 的 nn.Linear 和 nn.Conv2d 层默认使用 Kaiming 均匀初始化。

正交初始化

• nn.init.orthogonal_(tensor, gain=1.0): 用正交矩阵填充输入二维 tensor。这对于初始化循环神经网络 (RNN) 中的权重矩阵特别有用，有助于缓解循环连接中的梯度消失/爆炸问题。 Keras 提供 tf.keras.initializers.Orthogonal。

下图说明了 torch.nn.init 中的初始化函数如何应用于层的权重和偏置张量。

该图表明 torch.nn.init 中的函数直接就地修改层实例的权重和偏置张量。

在 PyTorch 模型中应用初始化

在 TensorFlow Keras 中，您通常在定义层时将初始化器指定为字符串标识符或初始化器对象，例如 kernel_initializer='glorot_uniform' 或 bias_initializer=tf.keras.initializers.Zeros()。

PyTorch 的方法更直接：您首先创建层实例，然后将 torch.nn.init 中的初始化函数应用于其权重和偏置张量。

初始化单个层

您可以访问层的参数（例如，nn.Linear 和 nn.Conv2d 的 weight 和 bias），并直接应用初始化函数：

import torch
import torch.nn as nn

# 示例：初始化一个线性层
linear_layer = nn.Linear(in_features=128, out_features=64)

# 对权重应用 Xavier 均匀初始化
nn.init.xavier_uniform_(linear_layer.weight)

# 如果偏置存在，将其初始化为零（nn.Linear 默认有偏置）
if linear_layer.bias is not None:
    nn.init.zeros_(linear_layer.bias)

print("已初始化 linear_layer.weight (前5个值):")
print(linear_layer.weight.data[0, :5])

这种直接操作对于精细控制或以不同方式初始化特定层很有用。

整个模型的自定义初始化函数

对于更复杂的模型，手动初始化每个层可能很繁琐。常见的做法是编写一个函数，它以模块 m 作为输入，检查其类型，并应用所需的初始化。然后，您可以使用 model.apply(your_init_function) 方法，该方法会递归地将函数应用于模型中的每个子模块。

import torch
import torch.nn as nn

# 定义一个示例模型
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.fc = nn.Linear(32 * 8 * 8, 10) # 假设输入图片尺寸导致 8x8 特征图

    def forward(self, x):
        x = self.relu1(self.conv1(x))
        x = self.relu2(self.conv2(x))
        x = x.view(x.size(0), -1) # 展平
        x = self.fc(x)
        return x

# 自定义权重初始化函数
def weights_init_custom(m):
    classname = m.__class__.__name__
    # 对于线性层
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.01) # 示例：偏置的小常量值
        print(f"Initialized {classname} with Xavier Uniform for weights and 0.01 for bias.")
    # 对于卷积层
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0)
        print(f"Initialized {classname} with Kaiming Normal for weights and 0 for bias.")
    # 您可以添加更多层类型，如 BatchNorm、RNN 等。

model = SimpleNet()
print("正在应用自定义权重初始化...")
model.apply(weights_init_custom)

# 您可以通过检查一些权重来验证
# print("\n自定义初始化后 conv1 的权重 (第一个过滤器，第一个通道):")
# print(model.conv1.weight.data[0, 0])

在 weights_init_custom 函数中：

• 我们使用 isinstance(m, nn.LayerType) 检查模块 m 的类型。
• 根据层类型，我们将所选的初始化应用于 m.weight 和 m.bias。
• model.apply(weights_init_custom) 调用确保此函数对 model 本身及其所有子模块（如 conv1、fc 等）执行。请注意，像 nn.ReLU 这样的激活层没有要初始化的参数，因此它们将被传递给函数，但 isinstance 检查将防止错误。

在 __init__ 方法中初始化

另一种简洁的管理初始化方式是直接在自定义 nn.Module 的 __init__ 方法中进行，就在定义每个层之后：

import torch
import torch.nn as nn

class ModelWithInit(nn.Module):
    def __init__(self, in_features, num_classes):
        super(ModelWithInit, self).__init__()
        self.linear1 = nn.Linear(in_features, 128)
        nn.init.kaiming_normal_(self.linear1.weight, nonlinearity='relu')
        if self.linear1.bias is not None:
            nn.init.zeros_(self.linear1.bias)
            
        self.relu = nn.ReLU()
        
        self.linear2 = nn.Linear(128, num_classes)
        nn.init.xavier_uniform_(self.linear2.weight) # 输出层示例
        if self.linear2.bias is not None:
            nn.init.zeros_(self.linear2.bias)

    def forward(self, x):
        x = self.relu(self.linear1(x))
        x = self.linear2(x)
        return x

# 创建模型实例
model_custom_init = ModelWithInit(in_features=784, num_classes=10)
# print("__init__ 中初始化后的 linear1 权重:")
# print(model_custom_init.linear1.weight.data[0, :5])

这种方式使初始化逻辑与层的定义紧密相关。

选择初始化器的考量

• 激活函数: 激活函数的选择是主要因素。
  • 对于后面接有 ReLU 或其变体的层，使用 Kaiming (He) 初始化。
  • 对于后面接有 sigmoid 或 tanh 的层，使用 Xavier (Glorot) 初始化。
• 网络架构: 更深的网络可能对初始化更敏感。
• 批量归一化: 像 nn.BatchNorm 这样的层可以使网络对初始权重尺度不那么敏感。然而，良好的初始化对于稳定性和速度仍然有利。
• 实验结果: 理论指导虽然有帮助，但有时针对特定问题或架构的最佳初始化方法是通过实验发现的。

通过了解并应用这些权重初始化技术，您可以显著提升 PyTorch 模型的稳定性和性能，借助您之前在 TensorFlow Keras 相关知识的基础上，同时适应 PyTorch 更明确的风格。