注意力机制的代码实现

一、前言

注意力机制的核心思想是，在处理一个序列时，模型应该关注序列中的哪些部分。在自然语言处理中，这通常意味着在生成一个词时，模型应该关注输入序列中的哪些词。注意力机制通过计算一个权重矩阵来实现这一点，这个权重矩阵表示输入序列中每个词对当前位置的重要性。

注意力机制的输入包括三个矩阵：query、key和value。这三个矩阵通常是通过嵌入层和神经网络层得到的。注意力机制的输出是一个加权后的value矩阵，这个矩阵表示了输入序列中每个词对当前位置的重要性。

以下是一些常见的注意力类型：

二、代码实现

解码器层需要把mask传入

def attention(query, key, value, mask=None, dropout=None):
    # 将query矩阵的最后一个维度值作为d_k
    d_k = query.size(-1)
    # 将key的最后两个维度互换(转置)，与query矩阵相乘，除以d_k开根号
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask==1, -1e9)
    p_attn = torch.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn