multiheadattention代码

以下是使用PyTorch实现的Multi-Head Attention模块的代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, n_heads, d_model):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.d_model = d_model
        
        # Linear layers for query, key and value projections
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        
        # Linear layer for final output
        self.fc = nn.Linear(d_model, d_model)
        
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Project query, key and value using linear layers
        Q = self.wq(query).view(batch_size, -1, self.n_heads, self.d_model // self.n_heads).transpose(1,2) # (batch_size, n_heads, seq_len, d_model // n_heads)
        K = self.wk(key).view(batch_size, -1, self.n_heads, self.d_model // self.n_heads).transpose(1,2) # (batch_size, n_heads, seq_len, d_model // n_heads)
        V = self.wv(value).view(batch_size, -1, self.n_heads, self.d_model // self.n_heads).transpose(1,2) # (batch_size, n_heads, seq_len, d_model // n_heads)
        
        # Compute scaled dot product attention for each head
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_model // self.n_heads, dtype=torch.float32)) # (batch_size, n_heads, seq_len, seq_len)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attention = F.softmax(scores, dim=-1)
        output = torch.matmul(attention, V) # (batch_size, n_heads, seq_len, d_model // n_heads)
        
        # Concatenate heads and apply final linear layer
        output = output.transpose(1,2).contiguous().view(batch_size, -1, self.n_heads * (self.d_model // self.n_heads)) # (batch_size, seq_len, d_model)
        output = self.fc(output)
        
        return output

该模块有三个输入：query，key和value，它们都是形状为(batch_size, seq_len, d_model)的张量。MultiHeadAttention首先使用三个线性层将query、key和value投影到d_model维空间，并将它们重塑为(batch_size, n_heads, seq_len, d_model // n_heads)的形状。然后，它对每个头计算缩放点积注意力，并将注意力权重和value相乘得到每个头的输出。最后，它将所有头的输出连接在一起并通过一个线性层输出最终的结果。

在forward方法中，我们首先计算query、key和value的投影。然后我们计算每个头的注意力权重，将它们与value相乘，并将每个头的输出连接在一起。最后，我们使用线性层将输出投影到d_model维空间。如果提供了掩码，则会将其应用于注意力权重。