写一段多头注意力机制的代码

好的，以下是一段使用 PyTorch 实现多头注意力机制的代码：

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        self.head_size = d_model // num_heads

        self.Wq = nn.Linear(d_model, d_model)
        self.Wk = nn.Linear(d_model, d_model)
        self.Wv = nn.Linear(d_model, d_model)

        self.fc = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        batch_size, seq_len, d_model = x.size()
        x = x.view(batch_size, seq_len, self.num_heads, self.head_size)
        return x.permute(0, 2, 1, 3)

    def forward(self, q, k, v, mask=None):
        q = self.Wq(q)
        k = self.Wk(k)
        v = self.Wv(v)

        q = self.split_heads(q)
        k = self.split_heads(k)
        v = self.split_heads(v)

        scores = torch.matmul(q, k.permute(0, 1, 3, 2))
        scores = scores / (self.d_model ** 0.5)

        if mask is not None:
            mask = mask.unsqueeze(1).unsqueeze(2)
            scores = scores.masked_fill(mask == 0, -1e9)

        attn_weights = nn.Softmax(dim=-1)(scores)
        attn_output = torch.matmul(attn_weights, v)

        attn_output = attn_output.permute(0, 2, 1, 3).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)

        attn_output = self.fc(attn_output)

        return attn_output, attn_weights

这个类实现了一个多头注意力机制的网络层，输入参数包括 d_model (输入向量维度) 和 num_heads (头的数量)。在 forward 方法中，我们首先将输入的查询、键和值向量分别通过线性层进行转换，并将它们分割成头。接下来，我们计算头之间的注意力得分，并将其乘以值向量。最后，我们将头拼接在一起，并通过全连接层进行后续处理。注意力权重也被返回，以便稍后进行可视化和分析。