手撕 Layer Normalization

它与 Batch Normalization 不同，Layer Normalization 是在单个样本的特征维度上进行归一化，而不是在批次维度上。具体步骤如下：

1. 计算均值和方差：对每个样本的所有特征计算均值和方差。

2. 标准化：将每个特征减去均值，除以标准差，使其分布的均值为 0，方差为 1。

3. 缩放和平移：通过可学习的参数 gamma 和 beta 对标准化后的结果进行缩放和平移。

import numpy
import torch.nn as nn


class my_LN(nn.Module):
    def __init__(self, embed_size, epsilon=1e-6):
        super(my_LN, self).__init__()
        self.gamma = nn.Parameter(torch.ones(embed_size))
        self.beta = nn.Parameter(zeros(embed_size))
        self.epsilon = epsilon

    def forward(self, x):
        # 1. 计算均值和方差
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)

        # 2. 标准化
        x_norm = (x-mean) / torch.sqrt(var + self.epsilon)

        # 3. 缩放和平移
        return self.gamma * x_norm + self.beta

# 设置随机种子以确保结果可复现
torch.manual_seed(42)

# 创建输入张量
batch_size, seq_len, embed_size = 2, 3, 4
input_tensor = torch.randn(batch_size, seq_len, embed_size)
print("输入张量 (input_tensor):")
print(input_tensor)

# 创建 LayerNorm 实例并应用
layer_norm = my_LN(embed_size=embed_size)
normalized_tensor = layer_norm(input_tensor)
print("\n经过 LayerNorm 处理后的张量 (normalized_tensor):")
print(normalized_tensor)