手撕经典算法 #2 神经网络篇

本文最后更新于：2025年4月10日 中午

本文对常见的几种神经网络组件进行了简单的实现和注释，便于理解。包括：

  • 层归一化（Layer Normalization，LN）
  • RMSNorm（Root Mean Square Layer Normalization）
  • 批次归一化（Batch Normalization，BN）
  • Dropout

LayerNorm

Layer Normalization (LN) 是一种归一化技术，旨在改善神经网络的训练稳定性和性能。LN 的基本思想是对每个样本在特征维度上进行归一化，而不是在批次维度上。对于每个输入 $ $ ，LN 的公式如下：

\[ \mu = \frac{1}{H} \sum_{i=1}^{H} x_i \]

\[ \sigma^2 = \frac{1}{H} \sum_{i=1}^{H} (x_i - \mu)^2 \]

\[ \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \]

其中，$ H $ 是特征维度的大小， $ $ 和 $ ^2 $ 分别是特征维度上的均值和方差。归一化后，会应用可学习的缩放参数 $ $ 和偏移参数 $ $ ：

\[ y_i = \gamma \hat{x}_i + \beta \]

代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import torch
from torch import nn

class LayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.hidden_size = hidden_size  # 隐藏状态的大小
        self.eps = eps  # 用于数值稳定性的一个小值
        
        # 初始化可学习的缩放和平移参数
        self.gamma = nn.Parameter(torch.ones(hidden_size))  # 缩放参数，初始值为全1
        self.beta = nn.Parameter(torch.zeros(hidden_size))  # 平移参数，初始值为全0
        
    def forward(self, x):
        # x 形状: (batch_size, seq_len, hidden_size)
        
        # 计算每个样本的均值和方差
        mean = x.mean(dim=-1, keepdim=True)  # 计算最后一个维度的均值，形状: (batch_size, seq_len, 1)
        variance = x.var(dim=-1, keepdim=True, unbiased=False)  # 计算最后一个维度的方差，形状: (batch_size, seq_len, 1)
        
        # 进行归一化
        x_normalized = (x - mean) / torch.sqrt(variance + self.eps)  # 归一化，形状: (batch_size, seq_len, hidden_size)
        
        # 应用缩放和平移参数
        output = self.gamma * x_normalized + self.beta  # 形状: (batch_size, seq_len, hidden_size)
        
        return output

def test_layer_norm():
    batch_size = 2
    seq_len = 4
    hidden_size = 8
    
    # 随机生成输入数据
    x = torch.randn(batch_size, seq_len, hidden_size)  # (batch_size, seq_len, hidden_size)
    
    # 创建 LayerNorm 模块
    layer_norm = LayerNorm(hidden_size)
    
    # 计算 LayerNorm 输出
    output = layer_norm(x)
    
    print("Input shape:", x.shape)
    print("Output shape:", output.shape)
    
if __name__ == "__main__":
	test_layer_norm()

RMSNorm

LayerNorm 需同时计算输入数据的均值和方差。研究发现 LayerNorm 的优势在于缩放不变性（即方差的计算，让模型对于输入和权重上的偏移噪声不敏感），而不是重新居中（即均值的计算，让模型在输入和权重都被随机缩放时保持输出表示不变）。

所以 RMSNorm 省略了归一化过程中的均值计算，仅计算特征的均方根（Root Mean Square），使得算法更加简洁，而效果不减，且运算效率显著提升。RMSNorm 的公式如下： \[ RMS=\sqrt{\frac{1}{H} \sum_{i=1}^{H} x_i^2 + \epsilon} \]

\[ \hat{x}_i = \frac{x_i}{RMS} \]

此外，通常仅保留缩放参数 $ $，省略偏移参数 $ $。实验表明，中心化（减均值）对性能影响有限，移除后仍能保持模型表达能力。

\[ y_i = \gamma \hat{x}_i \] 代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import torch
from torch import nn

class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.hidden_size = hidden_size
        self.eps = eps  # 用于数值稳定性的一个小值
        
        # 仅保留缩放参数 gamma，移除 beta
        self.gamma = nn.Parameter(torch.ones(hidden_size))  # 初始化为全 1
        
    def forward(self, x):
        # 计算均方根 (RMS)，不进行均值中心化
        rms = torch.sqrt(torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps)  # 平方后求均值，再开根
        
        # 归一化：x / RMS
        x_normalized = x / rms  # 不减去均值
        
        # 仅应用 gamma 参数，无 beta
        output = self.gamma * x_normalized
        
        return output

def test_rms_norm():  # 修改测试函数名
    batch_size = 2
    seq_len = 4
    hidden_size = 8
    
    x = torch.randn(batch_size, seq_len, hidden_size)
    rms_norm = RMSNorm(hidden_size)  # 使用 RMSNorm 类
    output = rms_norm(x)
    
    print("Input shape:", x.shape)
    print("Output shape:", output.shape)
    print("Params:", list(rms_norm.parameters()))  # 仅有一个参数 gamma
    
if __name__ == "__main__":
    test_rms_norm()  # 调用测试函数

BatchNorm

Batch Normalization (BN) 是另一种归一化技术，主要用于加速神经网络的训练。BN 对每个 mini-batch 的数据在批次维度上进行归一化。对于所有输入 $ $ ，BN 的公式如下：

\[ \mu_B = \frac{1}{m} \sum_{i=1}^{m} x_i \]

\[ \sigma_B^2 = \frac{1}{m} \sum_{i=1}^{m} (x_i - \mu_B)^2 \]

\[ \hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}} \]

其中， $ m $ 是 mini-batch 的大小， $ _B $ 和 $ _B^2 $ 分别是批次维度上的均值和方差。归一化后，也会应用可学习的缩放参数 $ $ 和偏移参数 $ $ ：

\[ y_i = \gamma \hat{x}_i + \beta \]

LN 和 BN 的区别：

  1. 归一化的维度：LN 在特征维度上进行归一化，而 BN 在批次维度上进行归一化。
  2. 应用场景：LN 更适用于 Recurrent Neural Networks (RNNs) 和 Transformer 等序列模型，而 BN 通常用于 Convolutional Neural Networks (CNNs)。
  3. 批量大小依赖性：LN 不依赖于批量大小，因此在小批量甚至单样本的情况下也能很好地工作。BN 依赖于较大的批量大小以稳定均值和方差的估计。

为什么 Transformer 使用 LN 而不是 BN：

  • Transformer 模型通常处理变长序列数据，其批次大小可能会变化或者在推理阶段可能只有一个样本。
  • BN 依赖于批次维度的均值和方差估计，因此在这种情况下表现可能不稳定。LN 则对每个样本独立进行归一化，不依赖于批次大小，因此更适合于 Transformer 这种模型。
  • 此外，LN 对序列模型的时间步长无关的归一化方式有助于保持输入数据的顺序特性，从而提高模型的性能和稳定性。

在 BatchNorm 的实现中，通常会区分训练（training）和推理（inference）阶段，这是因为在这两个阶段中，BN 的行为有所不同：

  1. 训练阶段（training）：

    • 在训练阶段，BN 会计算当前 mini-batch 的均值和方差，并使用这些统计量对数据进行归一化。具体公式同前文所述。
    • 训练过程中，BN 还会使用移动平均的方法更新运行时均值和方差，这样在推理阶段就可以使用这些全局统计量来代替 mini-batch 的统计量。更新规则如下： \[ \text{running\_mean} = (1 - \text{momentum}) \times \text{running\_mean} + \text{momentum} \times \mu_B \] \[ \text{running\_var} = (1 - \text{momentum}) \times \text{running\_var} + \text{momentum} \times \sigma_B^2 \]

  2. 推理阶段（inference）：

    • 在推理阶段，BN 不再使用 mini-batch 的均值和方差，而是使用训练阶段累积的运行时均值和方差。这是因为在推理阶段，通常批量大小很小甚至为 1，使用 mini-batch 的统计量会导致不稳定的输出。推理阶段的归一化公式如下： \[ \hat{x}_i = \frac{x_i - \text{running\_mean}}{\sqrt{\text{running\_var} + \epsilon}} \]

代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import torch
from torch import nn

class BatchNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-5, momentum=0.1):
        super().__init__()
        self.hidden_size = hidden_size  # 隐藏状态的大小
        self.eps = eps  # 用于数值稳定性的一个小值
        self.momentum = momentum  # 用于计算运行时均值和方差的动量
        
        # 初始化可学习的缩放和平移参数
        self.gamma = nn.Parameter(torch.ones(hidden_size))  # 缩放参数，初始值为全1
        self.beta = nn.Parameter(torch.zeros(hidden_size))  # 平移参数，初始值为全0
        
        # 初始化运行时均值和方差
        self.running_mean = torch.zeros(hidden_size)  # 运行时均值，初始值为全0
        self.running_var = torch.ones(hidden_size)  # 运行时方差，初始值为全1

    def forward(self, x):
        # x 形状: (batch_size, seq_len, hidden_size)
        if self.training:
            # 计算当前批次的均值和方差
            batch_mean = x.mean(dim=(0, 1), keepdim=False)  # 计算前两个维度的均值，形状: (hidden_size)
            batch_var = x.var(dim=(0, 1), keepdim=False, unbiased=False)  # 计算前两个维度的方差，形状: (hidden_size)

            # 更新运行时均值和方差
            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean
            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var
            
            mean = batch_mean
            variance = batch_var
        else:
            # 使用运行时均值和方差
            mean = self.running_mean
            variance = self.running_var
        
        # 进行归一化
        x_normalized = (x - mean) / torch.sqrt(variance + self.eps)  # 归一化，形状: (batch_size, seq_len, hidden_size)
        
        # 应用缩放和平移参数
        output = self.gamma * x_normalized + self.beta  # 形状: (batch_size, seq_len, hidden_size)
        
        return output

def test_batch_norm():
    batch_size = 2
    seq_len = 4
    hidden_size = 8
    
    # 随机生成输入数据
    x = torch.randn(batch_size, seq_len, hidden_size)  # (batch_size, seq_len, hidden_size)
    
    # 创建 BatchNorm 模块
    batch_norm = BatchNorm(hidden_size)
    
    # 计算 BatchNorm 输出
    output = batch_norm(x)
    
    print("Input shape:", x.shape)
    print("Output shape:", output.shape)
    
if __name__ == "__main__":
	test_batch_norm()

注意：这里的实现其实不太合理，因为 BN 通常不用于变长序列模型的输入，因此这里的 seq_len 维度只是摆设。

Dropout

Dropout 是一种正则化技术，用于防止神经网络的过拟合。通过在训练过程中随机「丢弃」一部分神经元，Dropout 使得模型不会过度依赖某些特定的神经元，从而增强模型的泛化能力。

在推理阶段，所有神经元都被激活，并根据 Dropout 概率进行缩放，主要目的是为了在训练和推理阶段保持一致的输出期望值。

具体而言，假设在训练阶段，输入神经元的激活值为 $ x $，Dropout 的概率为 $ p $。每个神经元以 $ 1 - p $ 的概率被保留（即不被丢弃）。因此，每个神经元的激活值期望为：

\[ E[\text{激活值}] = x \cdot (1 - p) \]

为了使得训练和推理阶段的输出期望一致，我们需要在训练阶段对保留的神经元进行缩放，即乘以 \(\frac{1}{1 - p}\)，这样可以抵消丢弃神经元带来的期望值减少。

代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import torch
from torch import nn

class Dropout(nn.Module):
    def __init__(self, dropout_prob=0.1):
        super().__init__()
        self.dropout_prob = dropout_prob  # Dropout 的概率
    
    def forward(self, x):
        if self.training:
            # 生成与输入形状相同的掩码，元素为 0 或 1，按照 dropout_prob 的概率为 0
            mask = (torch.rand(x.shape) > self.dropout_prob).float()  # 掩码，形状与 x 相同
            # 归一化掩码，使得训练阶段和推理阶段的一致性，除法保持训练和推理时的期望值一致
            output = mask * x / (1.0 - self.dropout_prob)  # 形状与 x 相同
        else:
            output = x  # 推理阶段，不进行 Dropout
        
        return output

def test_dropout():
    batch_size = 2
    seq_len = 4
    hidden_size = 8
    
    # 随机生成输入数据
    x = torch.randn(batch_size, seq_len, hidden_size)  # (batch_size, seq_len, hidden_size)
    
    # 创建 Dropout 模块
    dropout = Dropout(dropout_prob=0.1)
    
    # 设置为训练模式
    dropout.train()
    output_train = dropout(x)
    
    # 设置为推理模式
    dropout.eval()
    output_eval = dropout(x)
    
    print("Input shape:", x.shape)
    print("Output shape during training:", output_train.shape)
    print("Output shape during evaluation:", output_eval.shape)
    
if __name__ == "__main__":
	test_dropout()

Backpropagation

简单的示例代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch
import torch.nn as nn
import torch.optim as optim

# 加载必要组件，假设已经定义好
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
dataloader = ...

# 训练循环
model.train()
for epoch in range(num_epochs):
    for inputs, labels in data_loader:       
        # 前向传播
        preds = model(inputs)
        loss = criterion(preds, labels)
        
        # 反向传播与优化
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

Gradient Accumulation

Gradient Accumulation 核心思想是将多个小批量的梯度累积求和后，再统一更新模型参数，而非每个小批量单独更新。优点是在显存受限情况下，能够模拟更大 batch size 的训练。

梯度缩放：直接累积梯度作为一次参数更新的梯度，会使梯度值比实际大 gradient_accumulation_steps 倍。而Pytorch 的参数更新是写在optimizer.step()方法内部，无法手动控制。所以为了得到多个小批次的梯度平均后再更新模型，通常将 loss 除以 gradient_accumulation_steps。这一步等同于对 loss 进行缩放，来达到缩放梯度、求多个批次梯度的均值的目的。

梯度累加机制：在 PyTorch 中，每次反向传播 loss.backward()，计算得到的梯度会自动累积在 Tensor.grad 中，除非调用 optimizer.zero_grad() 手动清零。因此，只需控制好梯度清零与参数更新的间隔步数，无需重复累加操作。

代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
gradient_accumulation_steps = 4  

for batch_idx, (inputs, labels) in enumerate(data_loader):
    # 前向传播
    preds = model(inputs)
    loss  = criterion(preds, labels)

    # 关键，损失按累积步数缩放（等价于梯度求平均）
    loss = loss / gradient_accumulation_steps 

    # 反向传播，自动累加梯度
    loss.backward()			

    # 达到累积步数，或者处理末尾不完整批次时更新参数
    if ((batch_idx + 1) % gradient_accumulation_steps == 0) or (batch_idx + 1 == len(data_loader)):
        optimizer.step()
        optimizer.zero_grad()
算法笔记 > 手撕经典算法
#NLP #Algorithm
手撕经典算法 #2 神经网络篇
https://hwcoder.top/Manual-Coding-2
作者
Wei He
发布于
2024年7月7日
许可协议