手撕RMSNorm和LayerNorm

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-8):
        """
        d_model: 输入特征的维度
        eps: 防止除零的微小数值，确保数值稳定性
        """
        super().__init__()
        self.eps = eps
        # 可学习的缩放参数，初始化为1
        self.scale = nn.Parameter(torch.ones(d_model))
        # 可选：可学习的偏移参数，初始化为0
        self.bias = nn.Parameter(torch.zeros(d_model))

    def forward(self, x):
        # 计算RMS (均方根)，沿着最后一个维度进行
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        # 对输入进行RMS归一化
        x_norm = x / rms
        # 应用可学习的缩放参数和偏移
        return self.scale * x_norm + self.bias

# 示例
batch_size = 2
seq_len = 5
d_model = 10

# 随机生成输入张量
x = torch.randn(batch_size, seq_len, d_model)

# 创建RMSNorm层
rmsnorm = RMSNorm(d_model)

# 前向传播，输出归一化后的结果
output = rmsnorm(x)
print("输出张量：", output)

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, d_model, eps=1e-5):
        """
        d_model: 输入特征的维度
        eps: 防止除零的微小数值，确保数值稳定性
        """
        super().__init__()
        # 可学习的缩放参数，初始化为1
        self.scale = nn.Parameter(torch.ones(d_model))
        # 可学习的偏移参数，初始化为0
        self.bias = nn.Parameter(torch.zeros(d_model))
        self.eps = eps  # 小数值以防止除零错误

    def forward(self, x):
        # 计算均值和标准差，沿着最后一个维度（特征维度）进行
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)
        # 使用均值和标准差对输入进行归一化
        x_norm = (x - mean) / (std + self.eps)
        # 应用可学习的缩放参数和偏移参数
        return self.scale * x_norm + self.bias

# 示例
batch_size = 2
seq_len = 5
d_model = 10

# 随机生成输入张量
x = torch.randn(batch_size, seq_len, d_model)

# 创建 LayerNorm 层
layernorm = LayerNorm(d_model)

# 前向传播，输出归一化后的结果
output = layernorm(x)
print("输出张量：", output)