写一段代码pytorch实现self-attention

好的，下面是一个简单的 PyTorch 实现 self-attention 的代码：

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, hidden_size):
        super(SelfAttention, self).__init__()
        self.hidden_size = hidden_size
        self.query = nn.Linear(hidden_size, hidden_size)
        self.key = nn.Linear(hidden_size, hidden_size)
        self.value = nn.Linear(hidden_size, hidden_size)
        self.softmax = nn.Softmax(dim=2)
        
    def forward(self, x):
        query = self.query(x)
        key = self.key(x)
        value = self.value(x)
        scores = torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.hidden_size))
        scores = self.softmax(scores)
        output = torch.matmul(scores, value)
        
        return output

这个 self-attention 实现使用了三个线性层来计算查询、键和值。然后，通过矩阵乘法来计算得分，并通过 softmax 函数进行归一化。最后，将得分与值相乘得到输出。