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.projection = nn.Sequential(
            nn.Linear(hidden_size, 64),
            nn.ReLU(True),
            nn.Linear(64, 1)
        )
        
    def forward(self, encoder_outputs):
        # encoder_outputs shape: [batch_size, sequence_length, hidden_size]
        energy = self.projection(encoder_outputs)
        weights = torch.softmax(energy.squeeze(-1), dim=1)
        # weights shape: [batch_size, sequence_length]
        outputs = (encoder_outputs * weights.unsqueeze(-1)).sum(dim=1)
        # outputs shape: [batch_size, hidden_size]
        return outputs

在这个示例代码中，我们定义了一个SelfAttention类，它包含一个projection层，用于将encoder_outputs映射到注意力能量张量，并使用softmax函数将能量张量转换成权重张量。最后，我们将encoder_outputs与权重张量相乘并求和，得到self-attention的输出张量。

在forward函数中，我们首先通过projection层将encoder_outputs映射到能量张量，然后使用softmax函数将能量张量转换成权重张量。接着，我们将encoder_outputs与权重张量相乘并求和，得到self-attention的输出张量。最后，我们返回输出张量作为函数的输出。