pytorch 多头自注意力机制代码

下面是一个使用PyTorch实现多头自注意力机制的示例代码：

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, input_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.head_dim = input_dim // num_heads
        
        self.q_linear = nn.Linear(input_dim, input_dim)
        self.k_linear = nn.Linear(input_dim, input_dim)
        self.v_linear = nn.Linear(input_dim, input_dim)
        self.output_linear = nn.Linear(input_dim, input_dim)
        
    def forward(self, q, k, v):
        batch_size = q.size(0)
        
        q = self.q_linear(q).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_linear(k).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_linear(v).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        
        attention_scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim).float())
        attention_weights = torch.softmax(attention_scores, dim=-1)
        
        attended_values = torch.matmul(attention_weights, v)
        
        outputs = attended_values.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
        
        return self.output_linear(outputs)

这个代码定义了一个名为`MultiHeadAttention`的PyTorch模型类，它会根据输入的维度和头的数目构建多头自注意力机制。在`forward`函数中，输入的`q`、`k`和`v`是形状为`(batch_size, seq_len, input_dim)`的张量，表示查询、键和值。然后，使用线性变换将它们映射到多头注意力机制的维度。通过计算注意力分数和权重，以及对值进行加权求和，最终生成输出。

你可以根据需要调整模型的输入维度和头的数目，并在训练过程中使用该模型。希望对你有所帮助！如有其他问题，请随时提问。