Transformer(认真)

看不懂,好难啊!!!

详情请参考

【NLP】(task2)图解attention+transformer(代码讲解)_发现问题,并解决问题-CSDN博客​​​​​​​

一、RNN和Transformer的不同

二、从整体宏观来理解 Transformer

三、从细节来理解 Transformer

3.1 Transformer 的输入

3.2 Encoder(编码器)

3.3 Self-Attention 整体理解

四、Self-Attention 的细节

4.1 计算Query 向量,Key 向量,Value 向量

4.2 计算 Attention Score(注意力分数)

五、使用矩阵计算 Self-Attention

六、多头注意力机制(multi-head attention)

七、代码实现矩阵计算 Attention

7.1 使用矩阵实现多组注意力的并行计算(用pytorch自带的MultiheadAttention函数)

7.2 手动实现计算 Attention

7.3 关键代码

八、其他部分

8.1 使用位置编码来表示序列的顺序

8.2 Decoder和Cross Attention

8.3 最后的线性层和 Softmax 层

九、pytorch手写transformer

位置编码

三、多头注意力机制

3.1 拆开看多头注意力机制


一、RNN和Transformer的不同

(1)在 RNN 中,每一个 time step 的计算都依赖于上一个 time step 的输出,这就使得所有的 time step 必须串行化,无法并行计算,如下图所示。

在这里插入图片描述

 2)transformer模型的其中一个优点,就是使得模型训练过程能够并行计算。在Transformer 中,所有 time step 的数据,都是经过 Self Attention 计算,使得整个运算过程可以并行化计算。

在这里插入图片描述

二、从整体宏观来理解 Transformer

在这里插入图片描述

 首先,我们将整个模型视为黑盒。在机器翻译任务中,接收一种语言的句子作为输入,然后将其翻译成其他语言输出。

在这里插入图片描述

 中间部分的 Transformer 可以拆分为 2 部分:左边是编码部分(encoding component),右边是解码部分(decoding component)。

在这里插入图片描述

 其中编码部分是多层的编码器(Encoder)组成

同理,解码部分也是由多层的解码器

在这里插入图片描述

 在这里插入图片描述

encoder由多层编码器组成,每层编码器在结构上都是一样的,但不同层编码器的权重参数是不同的。每层编码器里面,主要由以下两部分组成

  • Self-Attention Layer
  • Feed Forward Neural Network(前馈神经网络,缩写为 FFNN)

在这里插入图片描述

 输入编码器的文本数据,首先会经过一个 Self Attention 层:这个层处理一个词的时候,不仅会使用这个词本身的信息,也会使用句子中其他词的信息

接下来,Self Attention 层的输出会经过前馈神经网络。

同理,解码器也具有这两层,但是这两层中间还插入了一个 Encoder-Decoder Attention 层——这个层能帮助解码器聚焦于输入句子的相关部分

三、从细节来理解 Transformer

3.1 Transformer 的输入

实际中向量一般是 256 或者 512 维。为了简化起见,这里将每个词的转换为一个 4 维的词向量。

那么整个输入的句子是一个向量列表,其中有 3 个词向量。

在实际中,每个句子的长度不一样,我们会取一个适当的值,作为向量列表的长度。

如果一个句子达不到这个长度,那么就填充全为 0 的词向量;如果句子超出这个长度,则做截断。

在这里插入图片描述

3.2 Encoder(编码器)

在这里插入图片描述

 编码器(Encoder)接收的输入都是一个向量列表,输出也是大小同样的向量列表,然后接着输入下一个编码器。

第一 个/层 编码器的输入是词向量,而后面的编码器的输入是上一个编码器的输出。下面,我们来看这个向量列表在编码器里面是如何流动的。

在这里插入图片描述

 每个单词转换成一个向量之后,进入self-attention层,每个位置的单词得到新向量,然后再输入FFN神经网络。

在这里插入图片描述

 每个位置的词都经过 Self Attention 层,得到的每个输出向量都单独经过前馈神经网络层,每个向量经过的前馈神经网络都是一样的

3.3 Self-Attention 整体理解

假设我们想要翻译的句子是:

The animal didn't cross the street because it was too tired

这个句子中的 it 是一个指代词,那么 it 指的是什么呢?它是指 animal 还是street?这个问题对人来说,是很简单的,但是对算法来说并不是那么容易。当模型在处理(翻译)it 的时候,Self Attention机制能够让模型把it和animal关联起来。

而 Transformer 使用Self Attention机制,会把其他单词的理解融入处理当前的单词。

在这里插入图片描述

四、Self-Attention 的细节

4.1 计算Query 向量,Key 向量,Value 向量

计算 Self Attention 的第 1 步是:对输入编码器的每个词向量,都创建 3 个向量,分别是:Query 向量,Key 向量,Value 向量。

这 3 个向量是词向量分别和 3 个矩阵相乘得到的,而这个3个矩阵是我们要学习的参数。

注意,这 3 个新得到的向量一般比原来的词向量的长度更小。

在这里插入图片描述

 上图中,有两个词向量:Thinking 的词向量 x1 和 Machines 的词向量 x2。以 x1 为例,X1 乘以 WQ 得到 q1

q1 就是 X1 对应的 Query 向量。同理,X1 乘以 WK 得到 k1,k1 是 X1 对应的 Key 向量;X1 乘以 WV 得到 v1,v1 是 X1 对应的 Value 向量。

Query 向量,Key 向量,Value 向量是什么含义呢?

其实它们就是 3 个向量,给它们加上一个名称,可以让我们更好地理解 Self-Attention 的计算过程和逻辑含义

4.2 计算 Attention Score(注意力分数)

第 2 步,是计算 Attention Score(注意力分数)。

假设我们现在计算第一个词 Thinking 的 Attention Score(注意力分数),需要根据 Thinking 这个词,对句子中的其他每个词都计算一个分数。

这些分数决定了我们在编码Thinking这个词时,需要对句子中其他位置的每个词放置多少的注意力。

这些分数,是通过计算对应的 Query 向量和其他位置的每个词的 Key 向量的点积,而得到的。

在这里插入图片描述

 

 第 4 步,接着把这些分数经过一个 Softmax 层,Softmax可以将分数归一化,这样使得分数都是正数并且加起来等于 1。

在这里插入图片描述

 这些分数决定了在编码当前位置(这里的例子是第一个位置)的词时,对所有位置的词分别有多少的注意力。

第 5 步,得到每个位置的分数后,将每个分数分别与每个 Value 向量相乘。这种做法背后的直觉理解就是:对于分数高的位置,相乘后的值就越大,我们把更多的注意力放到了它们身上;

第 6 步是把上一步得到的向量相加,就得到了 Self Attention 层在这个位置(这里的例子是第一个位置)的输出。

在这里插入图片描述

五、使用矩阵计算 Self-Attention

第一步是计算 Query,Key,Value 的矩阵。

在这里插入图片描述

 在这里插入图片描述

六、多头注意力机制(multi-head attention)

它扩展了模型关注不同位置的能力。在上面的例子中,第一个位置的输出 z1 包含了句子中其他每个位置的很小一部分信息,但 z1 可能主要是由第一个位置的信息决定的。当我们翻译句子:The animal didn’t cross the street because it was too tired时,我们想让机器知道其中的it指代的是什么。这时,多头注意力机制会有帮助。
在这里插入图片描述

 在多头注意力机制中,我们为每组注意力维护单独的 WQ, WK, WV 权重矩阵。将输入 X 和每组注意力的WQ, WK, WV 相乘,得到 8 组 Q, K, V 矩阵。
接着,我们把每组 K, Q, V 计算得到每组的 Z 矩阵,就得到 8 个 Z 矩阵。

在这里插入图片描述

 接下来就有点麻烦了,因为前馈神经网络层接收的是 1 个矩阵(其中每行的向量表示一个词),而不是 8 个矩阵。所以我们需要一种方法,把 8 个矩阵整合为一个矩阵。
怎么才能做到呢?我们把矩阵拼接起来,然后和另一个权重矩阵W Q相乘。

在这里插入图片描述

  1. 把 8 个矩阵 {Z0,Z1…,Z7} 拼接起来
  2. 把拼接后的矩阵和 WO 权重矩阵相乘
  3. 得到最终的矩阵 Z,这个矩阵包含了所有 attention heads(注意力头) 的信息。这个矩阵会输入到 FFNN (Feed Forward Neural Network)层。

在这里插入图片描述

 在这里插入图片描述

 在这里插入图片描述

七、代码实现矩阵计算 Attention

torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None)

PyTorch 提供了 MultiheadAttention 来实现 attention 的计算。

  • embed_dim:最终输出的 K、Q、V 矩阵的维度,这个维度需要和词向量的维度一样
  • num_heads:设置多头注意力的数量。如果设置为 1,那么只使用一组注意力。如果设置为其他数值,那么 - - num_heads 的值需要能够被 embed_dim 整除
  • dropout:这个 dropout 加在 attention score 后面

为什么 num_heads 的值需要能够被 embed_dim 整除:

这是为了把词的隐向量长度平分到每一组,这样多组注意力也能够放到一个矩阵里,从而并行计算多头注意力。

7.1 使用矩阵实现多组注意力的并行计算(用pytorch自带的MultiheadAttention函数)

在这里插入图片描述

 利用上面提的pytorch中的定义 MultiheadAttention 的对象后,调用时传入的参数如下

forward(query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None)

query:对应于 Key 矩阵,形状是 (L,N,E) 。其中 L 是输出序列长度,N 是 batch size,E 是词向量的维度

key:对应于 Key 矩阵,形状是 (S,N,E) 。其中 S 是输入序列长度,N 是 batch size,E 是词向量的维度

value:对应于 Value 矩阵,形状是 (S,N,E) 。其中 S 是输入序列长度,N 是 batch size,E 是词向量的维度


key_padding_mask:如果提供了这个参数,那么计算 attention score 时,忽略 Key 矩阵中某些 padding 元素,不参与计算 attention

  • 如果 key_padding_mask 是 ByteTensor,那么非 0 元素对应的位置会被忽略
  • 如果 key_padding_mask 是 BoolTensor,那么 True 对应的位置会被忽略

attn_mask:计算输出时,忽略某些位置。形状可以是 2D (L,S),或者 3D (N∗numheads,L,S)。其中 L 是输出序列长度,S 是输入序列长度,N 是 batch size。

如果 attn_mask 是 ByteTensor,那么非 0 元素对应的位置会被忽略
如果 attn_mask 是 BoolTensor,那么 True 对应的位置会被忽略

在这里插入图片描述

import torch
from torch import nn
## nn.MultiheadAttention 输入第0维为length
# batch_size 为 64,有 12 个词,每个词的 Query 向量是 300 维
query = torch.rand(12,64,300)
# batch_size 为 64,有 10 个词,每个词的 Key 向量是 300 维
key = torch.rand(10,64,300)
# batch_size 为 64,有 10 个词,每个词的 Value 向量是 300 维
value= torch.rand(10,64,300)

embed_dim = 300
num_heads = 10
# 输出是 (attn_output, attn_output_weights)
multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
attn_output = multihead_attn(query, key, value)[0]
# output: torch.Size([12, 64, 300])
# batch_size 为 64,有 12 个词,每个词的向量是 300 维
print(attn_output.shape)
# 打印出torch.Size([12, 64, 300])

 embed_dim是最终输出的 K、Q、V 矩阵的维度,这个维度需要和词向量的维度一样;多头注意力的数量num_heads值需要能够被 embed_dim 整除。

7.2 手动实现计算 Attention

在 PyTorch 提供的 MultiheadAttention 中,第 1 维是句子长度,第 2 维是 batch size。这里我们的代码实现中,第 1 维是 batch size,第 2 维是句子长度。

在这里插入图片描述

import torch
from torch import nn

class MultiheadAttention(nn.Module):
    # n_heads:多头注意力的数量
    # hid_dim:每个词输出的向量维度
    def __init__(self, hid_dim, n_heads, dropout):
        super(MultiheadAttention, self).__init__()
        self.hid_dim = hid_dim
        self.n_heads = n_heads

        # 强制 hid_dim 必须整除 h
        assert hid_dim % n_heads == 0
        # 定义 W_q 矩阵
        self.w_q = nn.Linear(hid_dim, hid_dim)
        # 定义 W_k 矩阵
        self.w_k = nn.Linear(hid_dim, hid_dim)
        # 定义 W_v 矩阵
        self.w_v = nn.Linear(hid_dim, hid_dim)
        self.fc = nn.Linear(hid_dim, hid_dim)
        self.do = nn.Dropout(dropout)
        # 缩放
        self.scale = torch.sqrt(torch.FloatTensor([hid_dim // n_heads]))

    def forward(self, query, key, value, mask=None):
        # K: [64,10,300], batch_size 为 64,有 12 个词,每个词的 Query 向量是 300 维
        # V: [64,10,300], batch_size 为 64,有 10 个词,每个词的 Query 向量是 300 维
        # Q: [64,12,300], batch_size 为 64,有 10 个词,每个词的 Query 向量是 300 维
        bsz = query.shape[0]
        Q = self.w_q(query)
        K = self.w_k(key)
        V = self.w_v(value)
        # 这里把 K Q V 矩阵拆分为多组注意力,变成了一个 4 维的矩阵
        # 最后一维就是是用 self.hid_dim // self.n_heads 来得到的,表示每组注意力的向量长度, 每个 head 的向量长度是:300/6=50
        # 64 表示 batch size,6 表示有 6组注意力,10 表示有 10 词,50 表示每组注意力的词的向量长度
        # K: [64,10,300] 拆分多组注意力 -> [64,10,6,50] 转置得到 -> [64,6,10,50]
        # V: [64,10,300] 拆分多组注意力 -> [64,10,6,50] 转置得到 -> [64,6,10,50]
        # Q: [64,12,300] 拆分多组注意力 -> [64,12,6,50] 转置得到 -> [64,6,12,50]
        # 转置是为了把注意力的数量 6 放到前面,把 10 和 50 放到后面,方便下面计算
        Q = Q.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        K = K.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        V = V.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)

        # 第 1 步:Q 乘以 K的转置,除以scale
        # [64,6,12,50] * [64,6,50,10] = [64,6,12,10]
        # attention:[64,6,12,10]
        attention = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale

        # 把 mask 不为空,那么就把 mask 为 0 的位置的 attention 分数设置为 -1e10
        if mask is not None:
            attention = attention.masked_fill(mask == 0, -1e10)

        # 第 2 步:计算上一步结果的 softmax,再经过 dropout,得到 attention。
        # 注意,这里是对最后一维做 softmax,也就是在输入序列的维度做 softmax
        # attention: [64,6,12,10]
        attention = self.do(torch.softmax(attention, dim=-1))

        # 第三步,attention结果与V相乘,得到多头注意力的结果
        # [64,6,12,10] * [64,6,10,50] = [64,6,12,50]
        # x: [64,6,12,50]
        x = torch.matmul(attention, V)

        # 因为 query 有 12 个词,所以把 12 放到前面,把 5 和 60 放到后面,方便下面拼接多组的结果
        # x: [64,6,12,50] 转置-> [64,12,6,50]
        x = x.permute(0, 2, 1, 3).contiguous()
        # 这里的矩阵转换就是:把多组注意力的结果拼接起来
        # 最终结果就是 [64,12,300]
        # x: [64,12,6,50] -> [64,12,300]
        x = x.view(bsz, -1, self.n_heads * (self.hid_dim // self.n_heads))
        x = self.fc(x)
        return x


# batch_size 为 64,有 12 个词,每个词的 Query 向量是 300 维
query = torch.rand(64, 12, 300)
# batch_size 为 64,有 12 个词,每个词的 Key 向量是 300 维
key = torch.rand(64, 10, 300)
# batch_size 为 64,有 10 个词,每个词的 Value 向量是 300 维
value = torch.rand(64, 10, 300)
attention = MultiheadAttention(hid_dim=300, n_heads=6, dropout=0.1)
output = attention(query, key, value)
## output: torch.Size([64, 12, 300])
print(output.shape)
# 打印出torch.Size([64, 12, 300])

7.3 关键代码

		# 这里把 K Q V 矩阵拆分为多组注意力,变成了一个 4 维的矩阵
        # 最后一维就是是用 self.hid_dim // self.n_heads 来得到的,表示每组注意力的向量长度, 每个 head 的向量长度是:300/6=50
        # 64 表示 batch size,6 表示有 6组注意力,10 表示有 10 个词,50 表示每组注意力的词的向量长度
        # K: [64,10,300] 拆分多组注意力 -> [64,10,6,50] 转置得到 -> [64,6,10,50]
        # V: [64,10,300] 拆分多组注意力 -> [64,10,6,50] 转置得到 -> [64,6,10,50]
        # Q: [64,12,300] 拆分多组注意力 -> [64,12,6,50] 转置得到 -> [64,6,12,50]
        # 转置是为了把注意力的数量 6 放到前面,把 10 和 50 放到后面,方便下面计算
        Q = Q.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        K = K.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        V = V.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)

经过 attention 计算得到 x 的形状是 `[64,12,6,50]`,64 表示 batch size,6 表示有 6组注意力,10 表示有 10 个词,50 表示每组注意力的词的向量长度。把这个矩阵转换为 `[64,12,300]`的矩阵,就是相当于把多组注意力的结果拼接起来。

这里的矩阵转换就是:把多组注意力的结果拼接起来,最终结果就是 [64,12,300],x: [64,12,6,50] -> [64,12,300]
x = x.view(bsz, -1, self.n_heads * (self.hid_dim // self.n_heads))

八、其他部分

8.1 使用位置编码来表示序列的顺序

目前为止,我们阐述的模型中缺失了一个东西,那就是表示序列中单词顺序的方法。

为了解决这个问题,Transformer 模型对每个输入的向量都添加了一个向量。这些向量遵循模型学习到的特定模式,有助于确定每个单词的位置,或者句子中不同单词之间的距离。这种做法背后的直觉是:将这些表示位置的向量添加到词向量中,得到了新的向量,这些新向量映射到 Q/K/V,然后计算点积得到 attention 时,可以提供有意义的信息。
在这里插入图片描述

 为了让模型了解单词的顺序,我们添加了带有位置编码的向量–这些向量的值遵循特定的模式。
如果我们假设词向量的维度是 4,那么带有位置编码的向量可能如下所示:

在这里插入图片描述

8.2 Decoder和Cross Attention

上面说了,编码器一般有多层,第一个编码器的输入是一个序列,最后一个编码器输出是一组注意力向量 K 和 V。这些注意力向量将会输入到每个解码器的Encoder-Decoder Attention层,这有助于解码器把注意力集中中输入序列的合适位置

在这里插入图片描述

 解码器中的 Self Attention 层,和编码器中的 Self Attention 层不太一样:在解码器里,Self Attention 层只允许关注到输出序列中早于当前位置之前的单词。具体做法是:在 Self Attention 分数经过 Softmax 层之前,屏蔽当前位置之后的那些位置。
在这里插入图片描述

8.3 最后的线性层和 Softmax 层

在这里插入图片描述

 线性层就是一个普通的全连接神经网络,可以把解码器输出的向量,映射到一个更长的向量,这个向量称为 logits 向量。

 现在假设我们的模型有 10000 个英语单词(模型的输出词汇表),这些单词是从训练集中学到的。因此 logits 向量有 10000 个数字,每个数表示一个单词的分数。我们就是这样去理解线性层的输出。然后,Softmax 层会把这些分数转换为概率

在这里插入图片描述

九、pytorch手写transformer

在这里插入图片描述

 所有需要的包的导入:

import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
from torch.nn.init import xavier_uniform_
from torch.nn.init import constant_
from torch.nn.init import xavier_normal_
import torch.nn.functional as F
from typing import Optional, Tuple, Any
from typing import List, Optional, Tuple
import math
import warnings
X = torch.zeros((2,4),dtype=torch.long)
embed = nn.Embedding(10,8)
print(embed(X).shape)
# 打印出torch.Size([2, 4, 8])

位置编码

Tensor = torch.Tensor
def positional_encoding(X, num_features, dropout_p=0.1, max_len=512) -> Tensor:
    r'''
        给输入加入位置编码
    参数:
        - num_features: 输入进来的维度
        - dropout_p: dropout的概率,当其为非零时执行dropout
        - max_len: 句子的最大长度,默认512
    
    形状:
        - 输入: [batch_size, seq_length, num_features]
        - 输出: [batch_size, seq_length, num_features]

    例子:
        >>> X = torch.randn((2,4,10))
        >>> X = positional_encoding(X, 10)
        >>> print(X.shape)
        >>> torch.Size([2, 4, 10])
    '''

    dropout = nn.Dropout(dropout_p)
    P = torch.zeros((1,max_len,num_features))
    X_ = torch.arange(max_len,dtype=torch.float32).reshape(-1,1) / torch.pow(
        10000,
        torch.arange(0,num_features,2,dtype=torch.float32) /num_features)
    P[:,:,0::2] = torch.sin(X_)
    P[:,:,1::2] = torch.cos(X_)
    X = X + P[:,:X.shape[1],:].to(X.device)
    return dropout(X)

# 位置编码例子
X = torch.randn((2,4,10))
X = positional_encoding(X, 10)
print(X.shape)
# 打印出torch.Size([2, 4, 10])

三、多头注意力机制

3.1 拆开看多头注意力机制

多头注意力类主要成分是:参数初始化、multi_head_attention_forward

if self._qkv_same_embed_dim is False:
    # 初始化前后形状维持不变
    # (seq_length x embed_dim) x (embed_dim x embed_dim) ==> (seq_length x embed_dim)
    self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim)))
    self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim)))
    self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim)))
    self.register_parameter('in_proj_weight', None)
else:
    self.in_proj_weight = Parameter(torch.empty((3 * embed_dim, embed_dim)))
    self.register_parameter('q_proj_weight', None)
    self.register_parameter('k_proj_weight', None)
    self.register_parameter('v_proj_weight', None)

if bias:
    self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
else:
    self.register_parameter('in_proj_bias', None)
# 后期会将所有头的注意力拼接在一起然后乘上权重矩阵输出
# out_proj是为了后期准备的
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
self._reset_parameters()

torch.empty是按照所给的形状形成对应的tensor,特点是填充的值还未初始化,类比torch.randn

在PyTorch中,变量类型是tensor的话是无法修改值的,而Parameter()函数可以看作为一种类型转变函数,将不可改值的tensor转换为可训练可修改的模型参数,即与model.parameters绑定在一起,register_parameter的意思是是否将这个参数放到model.parameters,None的意思是没有这个参数。
 

def _reset_parameters(self):
    if self._qkv_same_embed_dim:
        xavier_uniform_(self.in_proj_weight)
    else:
        xavier_uniform_(self.q_proj_weight)
        xavier_uniform_(self.k_proj_weight)
        xavier_uniform_(self.v_proj_weight)
    if self.in_proj_bias is not None:
        constant_(self.in_proj_bias, 0.)
        constant_(self.out_proj.bias, 0.)

B.multi_head_attention_forward

这个函数如下代码所示,主要分成3个部分:

  • query, key, value通过_in_projection_packed变换得到q,k,v
  • 遮挡机制
  • 点积注意力

import torch
Tensor = torch.Tensor
def multi_head_attention_forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    num_heads: int,
    in_proj_weight: Tensor,
    in_proj_bias: Optional[Tensor],
    dropout_p: float,
    out_proj_weight: Tensor,
    out_proj_bias: Optional[Tensor],
    training: bool = True,
    key_padding_mask: Optional[Tensor] = None,
    need_weights: bool = True,
    attn_mask: Optional[Tensor] = None,
    use_seperate_proj_weight = None,
    q_proj_weight: Optional[Tensor] = None,
    k_proj_weight: Optional[Tensor] = None,
    v_proj_weight: Optional[Tensor] = None,
) -> Tuple[Tensor, Optional[Tensor]]:
    r'''
    形状:
        输入:
        - query:`(L, N, E)`
        - key: `(S, N, E)`
        - value: `(S, N, E)`
        - key_padding_mask: `(N, S)`
        - attn_mask: `(L, S)` or `(N * num_heads, L, S)`
        输出:
        - attn_output:`(L, N, E)`
        - attn_output_weights:`(N, L, S)`
    '''
    tgt_len, bsz, embed_dim = query.shape
    src_len, _, _ = key.shape
    head_dim = embed_dim // num_heads
    q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)

    if attn_mask is not None:
        if attn_mask.dtype == torch.uint8:
            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
            attn_mask = attn_mask.to(torch.bool)
        else:
            assert attn_mask.is_floating_point() or attn_mask.dtype == torch.bool, \
                f"Only float, byte, and bool types are supported for attn_mask, not {attn_mask.dtype}"

        if attn_mask.dim() == 2:
            correct_2d_size = (tgt_len, src_len)
            if attn_mask.shape != correct_2d_size:
                raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
            attn_mask = attn_mask.unsqueeze(0)
        elif attn_mask.dim() == 3:
            correct_3d_size = (bsz * num_heads, tgt_len, src_len)
            if attn_mask.shape != correct_3d_size:
                raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
        else:
            raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")

    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
        warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
        key_padding_mask = key_padding_mask.to(torch.bool)
    
    # reshape q,k,v将Batch放在第一维以适合点积注意力
    # 同时为多头机制,将不同的头拼在一起组成一层
    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
    k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
    v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
    if key_padding_mask is not None:
        assert key_padding_mask.shape == (bsz, src_len), \
            f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
        key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).   \
            expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, src_len)
        if attn_mask is None:
            attn_mask = key_padding_mask
        elif attn_mask.dtype == torch.bool:
            attn_mask = attn_mask.logical_or(key_padding_mask)
        else:
            attn_mask = attn_mask.masked_fill(key_padding_mask, float("-inf"))
    # 若attn_mask值是布尔值,则将mask转换为float
    if attn_mask is not None and attn_mask.dtype == torch.bool:
        new_attn_mask = torch.zeros_like(attn_mask, dtype=torch.float)
        new_attn_mask.masked_fill_(attn_mask, float("-inf"))
        attn_mask = new_attn_mask

    # 若training为True时才应用dropout
    if not training:
        dropout_p = 0.0
    attn_output, attn_output_weights = _scaled_dot_product_attention(q, k, v, attn_mask, dropout_p)
    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
    attn_output = nn.functional.linear(attn_output, out_proj_weight, out_proj_bias)
    if need_weights:
        # average attention weights over heads
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        return attn_output, attn_output_weights.sum(dim=1) / num_heads
    else:
        return attn_output, None

  • 5
    点赞
  • 16
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值