多头注意机制pytorch代码

以下是一个简单的多头注意机制的 PyTorch 代码实现。假设我们有一个输入张量 x，维度为 (batch_size, seq_len, input_size)，我们希望得到一个输出张量 y，维度为 (batch_size, seq_len, output_size)，其中 output_size 是我们希望每个头注意力机制输出的向量大小。

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, input_size, output_size, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_heads = num_heads
        
        # 定义 Q、K、V 线性层
        self.q_linear = nn.Linear(input_size, output_size)
        self.k_linear = nn.Linear(input_size, output_size)
        self.v_linear = nn.Linear(input_size, output_size)
        
        # 定义输出线性层
        self.out = nn.Linear(output_size, output_size)
        
    def forward(self, x):
        bs = x.size(0)
        
        # 将输入张量 x 分别传入 Q、K、V 线性层
        q = self.q_linear(x).view(bs, -1, self.num_heads, self.output_size).transpose(1, 2) 
        k = self.k_linear(x).view(bs, -1, self.num_heads, self.output_size).transpose(1, 2)
        v = self.v_linear(x).view(bs, -1, self.num_heads, self.output_size).transpose(1, 2)
        
        # 计算 Q 和 K 的点积，并除以一个因子，以便更好地进行 softmax 操作
        scores = torch.matmul(q, k.transpose(-2, -1)) / self.output_size**0.5
        
        # 对 scores 进行 softmax 操作，得到注意力权重
        attention_weights = nn.functional.softmax(scores, dim=-1)
        
        # 将注意力权重和 V 相乘，得到多头注意力机制的输出
        out = torch.matmul(attention_weights, v).transpose(1, 2).contiguous().view(bs, -1, self.num_heads * self.output_size)
        
        # 将输出张量传入输出线性层，并返回结果
        out = self.out(out)
        return out

在这个代码中，我们首先定义了 MultiHeadAttention 类，该类包含了输入线性层（Q、K、V 线性层）和输出线性层。在 forward 函数中，我们首先将输入张量 x 传入 Q、K 和 V 线性层，再计算 Q 和 K 的点积，以获得注意力权重。然后，我们将注意力权重和 V 相乘，得到多头注意力机制的输出。最后，我们将输出张量传入输出线性层，并返回结果。

需要注意的是，在进行点积操作之前，我们将输出向量的维度除以一个因子，以便更好地进行 softmax 操作。这个因子是 $\sqrt{d}$，其中 $d$ 是输出向量的维度。在实现中，我们将这个因子保存在 self.output_size 中