import math import torch from torch import nn from d2l import torch as d2l def sequence_mask(X, valid_len, value=0): """在序列中屏蔽不相关的项 Defined in :numref:`sec_seq2seq_decoder`""" maxlen = X.size(1) #这里是被查询的单词的数,每个句子的长度是6 mask = torch.arange(maxlen, dtype=torch.float32, device=X.device)[None, :] < valid_len[:, None] X[~mask] = value return X def masked_softmax(X, valid_lens): """通过在最后一个轴上掩蔽元素来执行softmax操作 Defined in :numref:`sec_attention-scoring-functions`""" # X:3D张量,valid_lens:1D或2D张量 if valid_lens is None: return nn.functional.softmax(X, dim=-1) else: shape = X.shape if valid_lens.dim() == 1: valid_lens = torch.repeat_interleave(valid_lens, shape[1]) #2批次,5个头数,每句话4个单词。相当于每个单词5个查询头,所以相当于每批次每个句子的长度为10,但有效长度是3和2 else: valid_lens = valid_lens.reshape(-1) # 最后一轴上被掩蔽的元素使用一个非常大的负值替换,从而其softmax输出为0 X = sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value=-1e6) return nn.functional.softmax(X.reshape(shape), dim=-1) class DotProductAttention(nn.Module): """缩放点积注意力 Defined in :numref:`subsec_additive-attention`""" def __init__(self, dropout, **kwargs): super(DotProductAttention, self).__init__(**kwargs) self.dropout = nn.Dropout(dropout) # queries的形状:(batch_size,查询的个数,d) # keys的形状:(batch_size,“键-值”对的个数,d) # values的形状:(batch_size,“键-值”对的个数,值的维度) # valid_lens的形状:(batch_size,)或者(batch_size,查询的个数) def forward(self, queries, keys, values, valid_lens=None): d = queries.shape[-1] # 设置transpose_b=True为了交换keys的最后两个维度 scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d) self.attention_weights = masked_softmax(scores, valid_lens) return torch.bmm(self.dropout(self.attention_weights), values) #@save class MultiHeadAttention(nn.Module): """多头注意力""" def __init__(self, key_size, query_size, value_size, num_hiddens, num_heads, dropout, bias=False, **kwargs): super(MultiHeadAttention, self).__init__(**kwargs) self.num_heads = num_heads self.attention = DotProductAttention(dropout) self.W_q = nn.Linear(query_size, num_hiddens, bias=bias) self.W_k = nn.Linear(key_size, num_hiddens, bias=bias) self.W_v = nn.Linear(value_size, num_hiddens, bias=bias) self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias) def forward(self, queries, keys, values, valid_lens): # queries,keys,values的形状: # (batch_size,查询或者“键-值”对的个数,num_hiddens) # valid_lens 的形状: # (batch_size,)或(batch_size,查询的个数) # 经过变换后,输出的queries,keys,values 的形状: # (batch_size*num_heads,查询或者“键-值”对的个数,num_hiddens/num_heads) queries = transpose_qkv(self.W_q(queries), self.num_heads) #(10, 4, 20) keys = transpose_qkv(self.W_k(keys), self.num_heads) #(10, 6, 20) values = transpose_qkv(self.W_v(values), self.num_heads) #(10, 6, 20) if valid_lens is not None: # 在轴0,将第一项(标量或者矢量)复制num_heads次, # 然后如此复制第二项,然后诸如此类。 valid_lens = torch.repeat_interleave( #一个单词5个查询头,或者5个key,或者5个Value valid_lens, repeats=self.num_heads, dim=0) #这里用头数是因为每个单词有num_heads个头数,在masked_softmax中有单词数*头数的相似度分数需要设置为0 # output的形状:(batch_size*num_heads,查询的个数,num_hiddens/num_heads) output = self.attention(queries, keys, values, valid_lens) # output_concat的形状:(batch_size,查询的个数,num_hiddens) output_concat = transpose_output(output, self.num_heads) return self.W_o(output_concat) #@save def transpose_qkv(X, num_heads): """为了多注意力头的并行计算而变换形状""" # 输入X的形状:(batch_size,查询或者“键-值”对的个数,num_hiddens) # 输出X的形状:(batch_size,查询或者“键-值”对的个数,num_heads,num_hiddens/num_heads) X = X.reshape(X.shape[0], X.shape[1], num_heads, -1) # 输出X的形状:(batch_size,num_heads,查询或者“键-值”对的个数, num_hiddens/num_heads) X = X.permute(0, 2, 1, 3) # 最终输出的形状:(batch_size*num_heads,查询或者“键-值”对的个数, num_hiddens/num_heads) return X.reshape(-1, X.shape[2], X.shape[3]) # 第一个x进来,这里是2批次,5头,合计(10),查询的token数是4,每个变量的长度是20,所以查询变量数是10,token数是4,查询变量的向量长度是20 #@save def transpose_output(X, num_heads): """逆转transpose_qkv函数的操作""" X = X.reshape(-1, num_heads, X.shape[1], X.shape[2]) X = X.permute(0, 2, 1, 3) return X.reshape(X.shape[0], X.shape[1], -1) num_hiddens, num_heads = 100, 5 attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens, num_hiddens, num_heads, 0.5) attention.eval() batch_size, num_queries = 2, 4 num_kvpairs, valid_lens = 6, torch.tensor([3, 2]) X = torch.ones((batch_size, num_queries, num_hiddens))#(2, 4, 100)(批量、单词数、隐层输出后的单词向量大小) Y = torch.ones((batch_size, num_kvpairs, num_hiddens))#(2, 6, 100)(批量、单词数、隐层输出后的单词向量大小) print(attention(X, Y, Y, valid_lens).shape) # 查询4个单词和6个单词的多头注意力 """ 相当于单词长度变长了,需要掩码的查询变量变成了,单词*头数 头,头,头、头 头,头,头、头 ...... ^ ^ | | 单词----------- 单词--------单词------单词 """ """ 这个案例查询的单词数是4,即查询的句子的单词数是4,被查询的句子的单词数是6,5个头,即Q、K、V都是5,所以每个单词需要查询的Q*V是5次, 反过来,如果数字为PADDING的空值,需要掩住的注意力分数是单词数*Q数,或者说单词数*K数,或者说单词数*V数 """
李沐花书多头注意力机制代码详解
最新推荐文章于 2024-09-27 10:11:28 发布