手写多头注意力(MHA)的实现

import torch


import torch.nn as nn


import torch.nn.functional as F





class MultiHeadAttention(nn.Module):


    def __init__(self, embed_size, num_heads):


        super(MultiHeadAttention, self).__init__()


        self.embed_size = embed_size


        self.num_heads = num_heads


        self.head_dim = embed_size // num_heads





        assert (


            self.head_dim * num_heads == embed_size


        ), "Embedding size needs to be divisible by num_heads"





        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)


        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)


        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)


        self.fc_out = nn.Linear(num_heads * self.head_dim, embed_size)





    def forward(self, values, keys, query, mask):


        N = query.shape[0]


        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]





        # 将values, keys， queries分成多个头


        values = values.reshape(N, value_len, self.num_heads, self.head_dim)


        keys = keys.reshape(N, key_len, self.num_heads, self.head_dim)


        queries = query.reshape(N, query_len, self.num_heads, self.head_dim)


	   # 进行线性变换


        values = self.values(values)


        keys = self.keys(keys)


        queries = self.queries(queries)





        # Scaled dot-product attention（使用了爱因斯坦求和）


        # 这里直接求出kq矩阵


        attention = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])


        if mask is not None:


            attention = attention.masked_fill(mask == 0, float("-1e20"))


		# 使用softmax来归一化成注意力分数，分母是为了防止注意力分数相差过大


        attention = torch.softmax(attention / (self.embed_size ** (1 / 2)), dim=3)


		# 注意力分数和对应的value求加权和


        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(


            N, query_len, self.num_heads * self.head_dim


        )





        out = self.fc_out(out)


        return out








embed_size = 256


num_heads = 8


values = torch.randn(64, 10, embed_size)


keys = torch.randn(64, 10, embed_size)


query = torch.randn(64, 10, embed_size)


mask = None  # 可选的mask





multihead_attention = MultiHeadAttention(embed_size, num_heads)


output = multihead_attention(values, keys, query, mask)


print(output.shape)  # 期望输出: torch.Size([64, 10, 256])

实验结果：

Einsum 表示法是对张量的复杂操作的一种优雅方式，本质上是使用特定领域的语言。 一旦理解并掌握了 einsum，可以帮助我们更快地编写更简洁高效的代码。

基本内容：当两个变量具有相同的角标时，则遍历求和。在此情况下，求和号可以省略。

在实现一些算法时，数学表达式已经求出来了，需要将之转换为代码实现，使用einsum更加简单，但是可读性很差