Transformer的从0开始实现

在一个空间内，有query，keys，values，然后产生了注意力汇聚。

但是这样的空间只允许有一个么？可以有多个，然后最后统一输出么？
H1=f( W1(q) * Q, W1(k) * K, W1(v) * V)
H2=f( W2(q) * Q, W2(k) * K, W2(v) * V)
H3=f( W3(q) * Q, W3(k) * K, W3(v) * V)
…
Hi=f( Wi(q) * Q, Wi(k) * K, Wi(v) * V)

这里的f其实就是注意力评分函数，可以是加性的可以是点积。
最后再经由一个线性映射。
Output = W0 * ( H1, H2, H3, H4, H5….Hi)

但是理解不是最难的，这里比较难的点，是张量的分解和合并以此加强并行运算。

要用点积注意力评分，因为效率高。

import math
import torch
from torch import nn
from d2l import torch as d2l

def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # 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])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)
		
class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    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)

并行运算的张量分解和合并

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])


#@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)
		
测试：
#为了加强多个头的矩阵并行计算
#首先对于X，batch_size=2, 每个batch有3个query，每个query经过转换有num_hidden=8个特征。
#transpose_qky，会针对每个注意力头，有个batch，但是特征数目会除以head数目，比如2两个头，则每个query分出4个特征给头1，和头2.
#综上所述, 所以对于一个输入的无论是Q,K,V, 形状为(batch_size, 查询个数&kv对数目，特征数目)
#会转为(batch_size*num_head, 查询个数&kv对数目，特征数目/num_head)

X=torch.normal(0,1,(2,3,8))
print(X.shape)

Xt=transpose_qkv(X, 2)
print(Xt.shape)


#在进行完多头并行计算后，得到多个头的数据，这个时候需要合并多个头，进行最后的线性映射。
#(batch_size*num_head, 查询个数&kv对数目，特征数目/num_head) 通过transpose_output
Output=transpose_output(Xt, 2)
print(Output.shape)

结果：
torch.Size([2, 3, 8])
torch.Size([4, 3, 4])
torch.Size([2, 3, 8])

多头注意力

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=False)
        self.W_k=nn.Linear(key_size,num_hiddens,bias=False)
        self.W_v=nn.Linear(value_size,num_hiddens,bias=False)
        self.W_o=nn.Linear(num_hiddens,num_hiddens,bias=False)

    def forward(self,queries,keys,values,valid_lens):
        queriesT=transpose_qkv(self.W_q(queries), self.num_heads)
        keysT=transpose_qkv(self.W_k(keys), self.num_heads)
        valuesT=transpose_qkv(self.W_v(values), self.num_heads)
        
        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)
        output=self.attention(queriesT,keysT,valuesT,valid_lens)
        output_concat=transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

运行看看

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))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape

结果：
torch.Size([2, 4, 100])