import torch
import torch.nn as nn

# 设置随机种子保证实验可重复性
torch.manual_seed(123)

# 生成示例输入数据：batch_size=2, sequence_length=5（模拟两个样本的特征向量）
batch_example = torch.randn(2, 5)

# 关闭科学计数法显示，便于观察数值分布
torch.set_printoptions(sci_mode=False)

class LayerNorm(nn.Module):
    """层归一化实现（参考论文《Layer Normalization》核心思想）"""
    def __init__(self, emb_dim):
        """
        初始化层归一化模块
        :param emb_dim: 特征维度大小（对应输入张量的最后一维）
        """
        super().__init__()
        self.eps = 1e-5  # 数值稳定项，防止除以零（类似BN中的epsilon）
        self.scale = nn.Parameter(torch.ones(emb_dim))  # 可学习的缩放参数γ
        self.shift = nn.Parameter(torch.zeros(emb_dim)) # 可学习的平移参数β

    def forward(self, x):
        """
        前向传播过程（核心计算公式与原始论文一致）
        :param x: 输入张量，形状为(batch_size, embedding_dim)
        """
        # 计算沿特征维度的均值（保持维度用于广播）
        mean = x.mean(dim=-1, keepdim=True)
        # 计算沿特征维度的方差（使用有偏估计与PyTorch官方实现保持一致）
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        
        # 标准化计算：(x - μ)/√(σ² + ε)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        
        # 应用仿射变换：γ * x_normalized + β
        return self.scale * norm_x + self.shift

# 实例化层归一化模块（特征维度=5）
ln = LayerNorm(emb_dim=5)

# 前向传播计算
out_ln = ln(batch_example)

# 验证归一化效果
mean = out_ln.mean(dim=-1, keepdim=True)  # 沿特征维度计算均值
var = out_ln.var(dim=-1, unbiased=False, keepdim=True)  # 沿特征维度计算方差

print("归一化后均值:\n", mean)  # 理论值应接近0
print("归一化后方差:\n", var)   # 理论值应接近1

版权所有

版权归属：NateHHX

许可证：署名 4.0 国际 (CC-BY-4.0)