attention的pytorch示例代码如何写

以下是一个简单的PyTorch示例代码，用于实现self-attention机制：

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super(SelfAttention, self).__init__()
        
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.softmax = nn.Softmax(dim=-1)
        
    def forward(self, x):
        q = self.q_linear(x)
        k = self.k_linear(x)
        v = self.v_linear(x)
        
        attention_scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.embed_dim).float())
        attention_prob = self.softmax(attention_scores)
        attention_outputs = torch.matmul(attention_prob, v)
        
        return attention_outputs

在这个示例中，self-attention可以应用于输入张量x。首先，将x传递给三个不同的线性层，以生成密钥、查询和值。然后，将查询和密钥传递给一个点积注意力函数，以计算维度为（batch_size, seq_len, seq_len）的注意力矩阵。使用softmax函数作用于注意力矩阵以计算每个值的权重，并将其与值相乘以得到输出。请注意，我们还将注意力分数缩小了一个因子，以避免过度缩放。