首页 > 其他分享 >PyTorch中实现Transformer模型

PyTorch中实现Transformer模型

时间:2024-02-01 16:48:57浏览次数:29  
标签:__ src Transformer self mask PyTorch model dropout 模型

前言

  1. 关于Transformer原理与论文的介绍:详细了解Transformer:Attention Is All You Need

对于论文给出的模型架构,使用 PyTorch 分别实现各个部分。

引入的相关库函数:

import copy
import torch
import math
from torch import nn
from torch.nn.functional import log_softmax

# module: 需要深拷贝的模块
# n: 拷贝的次数
# return: 深拷贝后的模块列表
def clones(module, n: int) -> list:
    return [copy.deepcopy(module) for _ in range(n)]

1. 编码器与解码器堆叠

Encoder 编码器

编码器由 N 个相同的编码层堆叠而成,每个编码层含两个子层:多头注意力层和前馈网络层。每个子层后跟着一层,用于残差连接与标准化。

Add & Norm 残差连接和标准化

对于上一层的结果:\({\rm SubLayer}(x)\)与输出上一层的变量:\(x\)做残差连接并进行标准化:\({\rm LayerNorm}(x + {\rm Sublayer}(x))\)。

# 层标准化
class LayerNorm(nn.Module):
    # 设置 features 形状的张量作为可学习的参数,初始化
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        # 初始化两个参数,α为权重,β为偏置
        self.a_2 = nn.Parameter(torch.ones(features))  
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        # 计算最后一个维度的均值、方差
        mean = x.mean(-1, keepdim=True)  
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

# 子层残差连接
class SublayerConnection(nn.Module):
    # size: 参数矩阵的shape, 
    # dropout_prob: dropout概率
    def __init__(self, size, dropout_prob):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(p=dropout_prob)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))
  • nn.Dropout()初始化参数p表示训练时,以概率 p 将输入张量的一些元素归零,对于没有归零的元素将乘以\(\frac{1}{1-p}\)。
  • 输入为任意形状的张量,输出为与输入张量形状相同并经过处理的张量。[Source]

Multi-Head Attention 多头注意力层

计算点乘注意力:$ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V$

# q, k, v: 表示公式中的 Q, K, V
# mask: 当输入存在掩码时,将 mask 对应位置设置为负无穷
# dropout: dropout层
# return: 注意力层的输出,以及注意力权重
def attention(q, k, v, mask=None, dropout=None):
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    
    return torch.matmul(p_attn, v), p_attn 

# 多头注意力
class MultiHeadedAttention(nn.Module):
    # h: 多头注意力的头数
    # d_model: 嵌入词的维度
    def __init__(self, h, d_model, dropout_prob=0.1):
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout_prob)

    def forward(self, q, k, v, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1) # 相同的mask应用于所有的注意力头h
        batch_size = q.size(0)

        # 1) 执行线性变换,将 d_model 维度的 x 分割成 h 个 d_k 维度
        q, k, v = [
            # 通过 view 改变张量形状,并使用 transpose 方法交换张量维度
            lin(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (q, k, v))
        ]

        # 2) 将 attention 用于每个 batch 的投影向量上
        x, self.attn = attention(q, k, v, mask=mask, dropout=self.dropout)

        # 3) 通过线性层连接多头注意力计算完的向量
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
        return self.linears[-1](x)

关于contiguous()transpose()不改变张量物理上的存储顺序,而是改变了查看时逻辑上的顺序,使得在内存上不连续(可以通过is_contiguous()查看张量是否是连续的)。

如果不是连续的,可以通过contiguous()方法返回内存上连续、数值上相同的张量。view()方法改变张量的形状需要张量是连续的。[Source]

Feed Forward 前馈网络层

由两个线性层组成,中间使用 ReLU 激活函数:\(\mathrm{FFN}(x)=\max(0, xW_1 + b_1) W_2 + b_2\)

# 基于位置的前馈网络
class PositionwiseFeedForward(nn.Module):
    # d_model: 嵌入词的维度
    # d_ff: 前馈网络中间层的维度
    def __init__(self, d_model, d_ff, dropout_prob=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(p=dropout_prob)

    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))

编码层

每个编码层,含一个多头注意力层,一个前馈网络层,以及两个用于残差连接与标准化层分别跟在两个子层后面。N 个编码层组成编码器,每层的编码层的输出作为下一层的输入。

# 编码层
class EncoderLayer(nn.Module):
    # size: 参数矩阵的shape,
    # self_attn: 多头注意力层
    # feed_forward: 前馈网络层
    # dropout_prob: dropout概率
    def __init__(self, size, self_attn, feed_forward, dropout_prob):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout_prob), 2)
        self.size = size

    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda i: self.self_attn(i, i, i, mask))
        return self.sublayer[1](x, self.feed_forward)

# 编码器:由 N 个相同的层组成
class Encoder(nn.Module):
    def __init__(self, layer, n):
        super(Encoder, self).__init__()
        self.layers = clones(layer, n)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

EncoderLayerforward()内的x = self.sublayer[0](x, lambda i: self.self_attn(i, i, i, mask)),虽然此处输入的 q,k,v 均为 i 但在注意力层内,它们将分别与对应的 Q,K,V 矩阵(由线性层Linear实现)相乘,得到用于计算注意力的 q,k,v 。

Decoder 解码器

解码器由 N 层解码层组成。结构与编码层相似,由三个子层组成:带掩码的多头注意力层,多头注意力层和前馈网络层。每个子层后跟着一层,用于残差连接与标准化。

对于第二个子层,输入每一解码层的 K,V 为Encoder(第 N 层的编码层)的输出。为了区别输入Encoder和Decoder的嵌入词,分别用 src(Source,源) 和 tgt(Target,目标) 表示。

# 解码层:由多头注意力层、源-目标注意力层和前馈神经网络组成
class DecoderLayer(nn.Module):
    # size: 参数矩阵的shape,
    # self_attn: 多头注意力层
    # src_attn: 源-目标注意力层
    # feed_forward: 前馈网络层
    # dropout_prob: dropout概率
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout_prob):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout_prob), 3)

    # x: 解码曾输入
    # memory: 编码器的输出
    # src_mask: 源嵌入词掩码
    # tgt_mask: 目标嵌入词掩码
    # return: 解码层的输出
    def forward(self, x, memory, src_mask, tgt_mask):
        m = memory
        x = self.sublayer[0](x, lambda i: self.self_attn(i, i, i, tgt_mask))
        x = self.sublayer[1](x, lambda i: self.src_attn(i, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)


# 解码器:由 N 个相同的层组成
class Decoder(nn.Module):
    def __init__(self, layer, n):
        super(Decoder, self).__init__()
        self.layers = clones(layer, n)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

2. Generator 生成器

生成器将解码器的输出映射到词汇表上,由一个线性层和一个 softmax 层组成,用于预测下一个token的概率。

# 生成器:线性层和 softmax 层
class Generator(nn.Module):
    # d_model: 解码器输出的(嵌入词)向量维度
    # vocab: 词汇表的维度大小
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return log_softmax(self.proj(x), dim=-1)  # 对最后一个维度进行 softmax

3. Embedding 嵌入层

使用nn.Embedding构建查找表(Look-Up Table, LUT)。[Source]

  • 初始化时,num_embedding表示嵌入字典大小;embedding_dim表示每个嵌入词向量的维度大小。
  • forward()中使用时,输入维度为\(d\)的张量,返回维度为 \(d\times {\rm embedding\_dim}\) 的张量。

文中,作者还将嵌入层返回的张量乘以\(\sqrt{d_{model}}\)。

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(num_embeddings=vocab, embedding_dim=d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

4. Positional Encoding 位置编码

为了使模型学习文本的顺序信息,需要引入位置编码:

\[\begin{cases} PE_{(pos,2i)} = \sin(pos / 10000^{2i/d_{\text{model}}}) \\ PE_{(pos,2i+1)} = \cos(pos / 10000^{2i/d_{\text{model}}}) \end{cases} \]

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout_prob, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout_prob)

        # 计算位置编码
        pe = torch.zeros(max_len, d_model)  # Shape: max_len x d_model
        position = torch.arange(0, max_len).unsqueeze(1)  # Shape: max_len x 1
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000) / d_model))
        res = position * div_term  # Shape: max_len x d_model/2
        pe[:, 0::2] = torch.sin(res)
        pe[:, 1::2] = torch.cos(res)
        pe = pe.unsqueeze(0)  # Shape: 1 x max_len x d_model
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)].requires_grad_(False)
        return self.dropout(x)

self.register_buffer()用于将模型训练参数之外的变量注册加缓存,通过register_buffer()登记过的张量,会自动成为模型中的参数,随着模型移动(gpu/cpu)而移动,但是不会随着梯度进行更新。

在PyTorch中,对于梯度更新的需求,有着不同的张量定义方式[2]。

5. 整体架构

class EncoderDecoder(nn.Module):
    # encoder: 编码器
    # decoder: 解码器
    # src_embed: 源嵌入层
    # tgt_embed: 目标嵌入层
    # generator: 生成器
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    # src: 源语言句子
    # src_mask: 源语言句子掩码
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)  # 编码器

    # memory: 编码器的输出
    # src_mask: 源语言句子掩码
    # tgt: 目标语言句子
    # tgt_mask: 目标语言句子掩码
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

    def forward(self, src, tgt, src_mask, tgt_mask):
        memory = self.encode(src, src_mask)
        res_dec = self.decode(memory, src_mask, tgt, tgt_mask)
        return self.generator(res_dec)


# src_vocab: 源语言词典大小
# tgt_vocab: 目标语言词典大小
# n: 编码器和解码器的层数
# d_model: 嵌入词的维度
# d_ff: 前馈网络中间层的维度
# h: 多头注意力的头数
# dropout_prb: dropout概率
# return: Transformer 模型
def make_model(src_vocab, tgt_vocab, n=6, d_model=512, d_ff=2048, h=8, dropout_prb=0.1):
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout_prb)
    position = PositionalEncoding(d_model, dropout_prb)
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout_prb), n),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout_prb), n),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab),
    )
    # 初始化参数
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

参考文献

  1. The Annotated Transformer
  2. 实测!PyTorch 中 nn.Parameter、register_buffer 和直接把 Tensor 当属性有啥区别?

标签:__,src,Transformer,self,mask,PyTorch,model,dropout,模型
From: https://www.cnblogs.com/zh-jp/p/18001551

相关文章

  • 手把手教你如何创建并上传modelscope模型
    参考来源:https://modelscope.cn/docs/模型的创建与文件上传1.注册modelscope相关账号(略)2.创建对应的模型3.填写模型的相关资料4.创建审核通过了之后,下载对应的模型文件夹5.拷贝对应的上传脚本,可以根据上面的页面复制使用modelscope的SDK脚本6.需要获取用户特......
  • 如何将PyTorch模型迁移到昇腾平台
    https://bbs.huaweicloud.com/blogs/399602?utm_source=cnblog&utm_medium=bbs-ex&utm_campaign=other&utm_content=content如何将PyTorch模型迁移到昇腾平台举报 昇腾CANN 发表于2023/04/1809:54:50  5k+  0  1 【摘要】本文介绍将PyTorch网络模型迁移到昇......
  • 4个大语言模型训练中的典型开源数据集
    本文分享自华为云社区《浅谈如何处理大语言模型训练数据之三开源数据集介绍》,作者:码上开花_Lancer。随着最近这些年来基于统计机器学习的自然语言处理的算法的发展,以及信息检索研究的需求,特别是近年来深度学习和预训练语言模型的研究以及国内国外许多大模型的开源,研究人员们构建......
  • 读论文-基于用户长短期偏好的序列推荐模型
    前言今天要读的论文名为《基于用户长短期偏好的序列推荐模型》,是一篇于2022年12月29日发表在《计算机科学》上的一篇期刊论文。文章发现了传统的序列推荐模型忽略了不同用户的个性化行为,导致模型不能充分捕获用户动态偏好而产生的兴趣漂移等问题,提出了一种基于用户长短期偏好......
  • 想要成为AIGC大模型工程师, 如何搭建你的知识体系框架?
    AI不会取代你的工作,会取代你的是会AI的人,如何提升自己与他人的知识壁垒,如何学习AIGC?相信在过去2023这个AI爆发的元年,我相信也是很多人的疑问,好像不懂使用点AI工具,例如ChatGpt、Midjourney或者SD,就好像被淘汰了一样,更近一步地我们如何训练自己的AI大模型、如何开发自己的AI产品?是很......
  • R语言GAMLSS模型对艾滋病病例、降雪量数据拟合、预测、置信区间实例可视化
    全文链接:http://tecdat.cn/?p=31996原文出处:拓端数据部落公众号GAMLSS模型是一种半参数回归模型,参数性体现在需要对响应变量作参数化分布的假设,非参数性体现在模型中解释变量的函数可以涉及非参数平滑函数,非参数平滑函数不预先设定函数关系,各个解释变量的非线性影响结果完全取决......
  • 【极简】Pytorch中的register_buffer()
    registerbuffer定义模型能用torch.save保存的、但是不更新参数。使用:只要是nn.Module的子类就能直接self.调用使用:classA(nn.Module):#...self.register_buffer('betas',torch.linspace(beta_1,beta_T,T).double())#...手动定义参数上述的参数显然可以......
  • AIGC、AGI、ChatGPT:揭秘人工智能大模型的变革与创新
    随着科技的不断发展,人工智能(AI)已经深入到我们生活的方方面面,从语音助手到自动驾驶汽车,再到智能家居,AI的影子无处不在。而在近年来,AIGC、AGI和ChatGPT等人工智能大模型的出现,更是引领了一场前所未有的技术革命。接下来,我们将一起探索这些人工智能大模型的奥秘。一、AIGC:开启全新的......
  • 大模型训练、多模态数据处理与融合:从理论到实践
    一、大模型训练大模型是指具有巨大参数量和计算能力的人工神经网络模型,如GPT(GenerativePre-trainedTransformer)系列模型和BERT(BidirectionalEncoderRepresentationsfromTransformers)模型。大模型的训练需要大量的数据和计算资源,同时需要采用合适的优化算法和技术,以提高模型的......
  • 模型评估与轨道
    模型评估与轨道一、模型评估的基本方法1.1监督学习下的泛化、过拟合与欠拟合在有监督的学习过程中,首先在训练数据上学得模型参数来构建模型,然后根据学得的模型,对新数据做出预测。用来训练的数据集称为训练集,用来测试预测结果是否准确的新数据称为测试集。注意:测试集中的数据不......