多头注意力

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


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 = d2l.DotProductAttention(dropout)  #点乘
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)  #100，100
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)   #100，100
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)   #100，100
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)  #100，100

    def forward(self, queries, keys, values, valid_lens):     #进行计算每一头的qkv
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)   #在轴0，将第一项复制5次，然后如此复制第二项，然后诸如此类。

        output = self.attention(queries, keys, values, valid_lens)  #<q,k>v 进行点乘

        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    X = X.permute(0, 2, 1, 3)

    return X.reshape(-1, X.shape[2], X.shape[3])  #最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,num_hiddens/num_heads)

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  #隐藏层100，5头注意力
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5) #将qkv的长度信息，还有网络层的一些设置输入多头注意力网络
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))
OUT = attention(X, Y, Y, valid_lens)  #做为qkv输入网络，进行变换
print(OUT)