使用 nn.MultiheadAttention实现self-attention

好的，下面是使用 `nn.MultiheadAttention` 实现 self-attention 的代码：

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super(SelfAttention, self).__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.multihead_attn = nn.MultiheadAttention(hidden_size, num_heads)
        self.layer_norm = nn.LayerNorm(hidden_size)
        
    def forward(self, x):
        attn_output, _ = self.multihead_attn(x, x, x)
        output = self.layer_norm(x + attn_output)
        
        return output

这个 self-attention 实现使用了 `nn.MultiheadAttention` 层来计算自注意力。该层需要输入三个相同的张量，分别为查询、键和值。我们将输入 `x` 作为所有三个张量的输入。输出是经过自注意力计算后的张量。最后，我们将层归一化应用于输入和输出的和。