【大模型手撕】pytorch实现LayerNorm, RMSNorm

LayerNorm介绍请参考:【AI知识】归一化、批量归一化 、 层归一化 和 实例归一化

RMSNorm介绍请参考:【大模型知识点】RMSNorm(Root Mean Square Normalization)均方根归一化

LayerNorm实现:

python 复制代码
import torch 
import torch.nn as nn


class LayerNorm(nn.Module):
    def __init__(self,dim,eps=1e-5,bias=False):
        super().__init__()
        self.dim = dim
        self.eps = eps
        # 可训练的缩放参数
        self.gamma = nn.Parameter(torch.ones(dim))

        self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
    
    def forward(self,x):
        # x: (batch_size,seq_len,dim)
        # 计算均值 x_mean : (batch_size,seq_len,dim)
        x_mean = x.mean(-1,keepdim=True)
        # 计算均方根 rms :  (batch_size,seq_len,dim)
        rms = torch.sqrt(x.pow(2).mean(-1,keepdim=True)+self.eps)

        if self.bias:
            return self.gamma*((x-x_mean)/rms)+self.bias
        else:
            return self.gamma*((x-x_mean)/rms)

RMSNorm实现:

python 复制代码
import torch 
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self,dim,eps=1e-5,bias=False):
   		super().__init__()
        self.dim = dim 
        self.eps = eps
        # 可训练的缩放参数
        self.gamma = nn.Parameter(torch.ones(dim))
        self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
    def forward(self,x):
        # 计算输入的均方根
        # x: (batch_size,seq_len,dim)
        # .mean(-1,keepdim=True) : 在最后一个维度(特征维度)上计算平均值,并保持维度不变
        # rms : (batch_size,seq_len,1)
        rms = torch.sqrt(x.pow(2).mean(-1,keepdim=True)+self.eps)

        if self.bias:
            return self.gamma*(x/rms) + self.bias
        else:
            return self.gamma*(x/rms)
相关推荐
岁月宁静18 小时前
深度定制:在 Vue 3.5 应用中集成流式 AI 写作助手的实践
前端·vue.js·人工智能
Dfreedom.18 小时前
一文掌握Python四大核心数据结构:变量、结构体、类与枚举
开发语言·数据结构·python·变量·数据类型
一半烟火以谋生18 小时前
Python + Pytest + Allure 自动化测试报告教程
开发语言·python·pytest
galaxylove18 小时前
Gartner发布数据安全态势管理市场指南:将功能扩展到AI的特定数据安全保护是DSPM发展方向
大数据·人工智能
格林威18 小时前
偏振相机在半导体制造的领域的应用
人工智能·深度学习·数码相机·计算机视觉·视觉检测·制造
叶子丶苏19 小时前
第八节_PySide6基本窗口控件_按钮类控件(QAbstractButton)
python·pyqt
晓枫-迷麟19 小时前
【文献阅读】当代MOF与机器学习
人工智能·机器学习
百锦再19 小时前
对前后端分离与前后端不分离(通常指服务端渲染)的架构进行全方位的对比分析
java·开发语言·python·架构·eclipse·php·maven
来酱何人19 小时前
实时NLP数据处理:流数据的清洗、特征提取与模型推理适配
人工智能·深度学习·分类·nlp·bert
sensen_kiss19 小时前
INT301 Bio-computation 生物计算(神经网络)Pt.3 梯度下降与Sigmoid激活函数
人工智能·神经网络·机器学习