Transformer代码学习—手写编码模块，encoder layer (网课)

Antonin_Artaud

已于 2024-06-16 17:11:26 修改

阅读量797

点赞数 17

文章标签： transformer 学习深度学习

于 2024-06-13 22:11:14 首次发布

本文链接：https://blog.csdn.net/mpu2050/article/details/139665330

版权

原作者视频链接：

【研1基本功（真的很简单）注意力机制】手写多头注意力机制_哔哩哔哩_bilibili先看文档哈 https://dwexzknzsh8.feishu.cn/docx/VkYud3H0zoDTrrxNX5lce0S4nDh?from=from_copylink慢慢更新，一周更完transformer同步更新RLHF、LoRA, 视频播放量 77272、弹幕量 52、点赞数 1996、投硬币枚数 1202、收藏人数 7088、转发人数 391, 视频作者 happy魇, 作者简介 Enjoy life | 欢迎来到我的主页！，相关视频：注意力机制的本质|Self-Attention|Transformer|QKV矩阵，【深度学习缝合模块】废材研究生自救指南！12个最新模块缝合模块创新！-CV、注意力机制、SE模块，膜拜！浙大教授竟把Transformer讲的如此简单！全套【Transformer基础】课程分享，连草履虫都能学会！再学不会UP下跪！，【可视化】Transformer中多头注意力的计算过程，24种魔改注意力机制暴力涨点即插即用 CNN+注意力机制，Transformer论文逐段精读【论文精读】，【官方双语】直观解释注意力机制，Transformer的核心 | 【深度学习第6章】，【研1基本功（真的很简单）Group Query-Attention】大模型训练必备方法——bonus(位置编码讲解)，视觉十分钟｜通道注意力原理（SENet，CBAM，SRM，ECA，FcaNet）｜计算机视觉/通道注意力/总结分享，图神经网络改进-手把手教你改代码-第1期https://www.bilibili.com/video/BV1o2421A7Dr/

1.参数初始化

import torch
from torch import nn
import torch.functional as f
import math
#%%
# 测试数据
X = torch.randn( 128, 64, 512) # Batch,Time,Dimension
print(X.shape)

2. 多头注意力完善(加入mask判断便于改为自注意力)

# 设置multihead_attention基本参数
d_model = 512  # 映射到Q，K，V空间中有多少位
n_head = 8 # 有多少个头
#%%
# 对上次实现的multihead attention进行修复（因为之前都是默认生成了一个mask）
# 但是在encoder当中做的是self_attention，不需要mask
# 因此需要我们多新建一个参数，来告诉他你是否需要mask
class multi_head_attention(nn.Module):
    def __init__(self, d_model,n_head) -> None:
        super(multi_head_attention,self).__init__()

        self.n_head = n_head
        self.d_model = d_model
        self.w_q = nn.Linear(d_model, d_model)  # 线性层映射函数，把初始向量映射到Q,K,V（query,key,value）
        # 简单来说就是去寻找一些query去跟key，问他（key）哪些数据是跟我匹配的上的，匹配上之后，key所对应的value值进行加权组合，最终得到attention的输出
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_combine = nn.Linear(d_model, d_model) # 由于是多头注意力，所以要在最后做一个组合映射(多写一个w_combine的线性映射)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, q, k, v, mask=None):
        batch, time, dimension = q.shape
        n_d = self.d_model // self.n_head   # 得到新维度
        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)  # 把qkv分别丢到上面定义的三个线性映射层中，就可以得到qkv空间中的一个表示
        # 对空间表示进行切分，对我们需要得到几个头进行切分
        q = q.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3) # 把q进行维度划分，一维是batch，二维是time， 三维是n.head（分成几个头），四维是n.d（分完头之后的维度）
        k = k.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3) # 也可以说把最后一维拆成了n.head和n.d两个维度的乘积
        v = v.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3) # 做attention操作的时候head维是不能放在最后的，对最后两个维度进行处理，所以要用permute指令做一个维度变换
        # 原先的维度是0，1，2，3现在则是0，2，1，3

        score = q @ k.transpose(2, 3) / math.sqrt(n_d) # q乘以k的转置除以它的维度开根号（让方差变小） @是矩阵乘法
        # torch.tril命令-生成下三角矩阵（左下角都是1，右上角都是0）

        if mask is not None:
            mask = torch.tril(torch.ones(time, time, dtype=bool))
            score = score.masked_fill(mask == 0, float("-inf")) # 把mask等于0的地方都填充为负无穷

        # 填充为负无穷的原因:softmax操作时e^-inf就是0，就相当于我们不去care后面部分的信息
        score = self.softmax(score) @ v

        # 最后把得分的格式变回来（因为之前把time维和self.n_head维度进行了旋转，现在则是要旋转回来），然后再过一个连续性函数
        score = score.permute(0, 2, 1, 3).contiguous().view(batch, time, dimension)
        # contiguous()的作用是让整个矩阵序列在内存中都是连续的

        output = self.w_combine(score)
        return output

attention = multi_head_attention(d_model, n_head)
output = attention(X, X, X)
print(output, output.shape)

3. Embedding

3.1 Token Embedding

class TokenEmbedding(nn.Embedding):
    def __init__(self, vocab_size, d_model):
        super(TokenEmbedding,self).__init__(vocab_size, d_model, padding_idx=1)

3.2 Position Embedding

class PositionalEmbedding(nn.Module):
    def __init__(self, d_model, maxlen, device):
        super(PositionalEmbedding, self).__init__()
        self.encoding = torch.zeros(maxlen, d_model, device) # 初始化编码（长度是maxlen）
        self.encoding.requires_grad_(False) # 告诉模型不用梯度

        pos = torch.arange(0, maxlen, device)
        pos = pos.float().unsqueeze(1) # 增加一个维度
        _2i = torch.arange(0, d_model, 2, device)

        self.encoding[:,0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:,1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))

    def forward(self, x):
        seq_len = x.shape[1]
        return self.encoding[:seq_len,:]

3.3 Total Embedding (综合前两者)

class TransformerEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, max_len, drop_prob, device):
        super(TransformerEmbedding, self).__init__()
        self.tok_emb = TokenEmbedding(vocab_size, d_model)
        self.pos_emb = PositionalEmbedding(d_model, max_len, device)
        self.drop_out = nn.Dropout(p=drop_prob)

    def forward(self, x):
        tok_emb = self.tok_emb(x)
        pos_emb = self.pos_emb(x)
        return self.drop_out(tok_emb + pos_emb)

4. LayerNorm (归一化)

# 归一化————学习两个参数-减均值除方差
class LayerNorm(nn.Module):
    def __init__(self, d_model, eps = 1e-10):
        super(LayerNorm, self).__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        var = x.var(-1, unbiased=False, keepdim=True)
        out = (x - mean) / torch.sqrt(var + self.eps)
        out = self.gamma * out + self.beta
        return out

5. FFN——feedforward层，本质上是一个两层的MLP，第一层会将输入的向量升维，第二层将向量重新降维，这样就可以学习到更加抽象的特征

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, hidden, dropout=0.1):
        self.fc1 = nn.Linear(d_model, hidden)
        self.fc2 = nn.Linear(hidden, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

6. Encoder layer

class EncoderLayer(nn.Module):
    # ffn_hidden为隐藏层维度，drop_prob为dropout的概率
    def __init__(self, d_model, ffn_hidden, n_head, drop_prob) -> None:
        super(EncoderLayer, self).__init__()
        self.attention = multi_head_attention(d_model, n_head)
        self.norm1 = LayerNorm(d_model)
        self.drop1 = nn.Dropout(drop_prob)

        self.ffn = PositionwiseFeedForward(d_model, ffn_hidden, drop_prob)
        self.norm2 = LayerNorm(d_model)
        self.drop2 = nn.Dropout(drop_prob)

    def forward(self, x, mask=None):
        # 因为等会要用到残差连接，所以要先生成一个备份
        _x = x
        x = self.attention(x, x, x, mask)

        x = self.drop1(x)
        self.norm1(x + _x)

        _x = x
        x = self.ffn(x)

        x = self.drop2(x)
        x = self.norm2(x + _x)
        return x