pytorch编程实现自注意力机制

好的，以下是一个简单的 PyTorch 实现自注意力机制的例子：

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        
        # 定义三个可学习的参数矩阵
        self.query_weight = nn.Linear(embed_dim, embed_dim, bias=False)
        self.key_weight = nn.Linear(embed_dim, embed_dim, bias=False)
        self.value_weight = nn.Linear(embed_dim, embed_dim, bias=False)

    def forward(self, x):
        batch_size, seq_len, embed_dim = x.size()
        # 将输入张量 x 沿着最后一维分成 num_heads 份
        x = x.view(batch_size, seq_len, self.num_heads, self.head_dim)
        # 将 x 转置后合并前两个维度，得到新的维度 (batch_size * num_heads, seq_len, head_dim)
        x = x.transpose(1, 2).contiguous().view(batch_size * self.num_heads, seq_len, self.head_dim)
        
        # 分别计算 query、key、value
        query = self.query_weight(x)
        key = self.key_weight(x)
        value = self.value_weight(x)
        
        # 将 query 和 key 进行点积，然后除以一个缩放因子 sqrt(d_k)
        scores = torch.bmm(query, key.transpose(1, 2)) / (self.head_dim ** 0.5)
        # 对 scores 进行 softmax 操作，得到注意力权重
        attn_weights = F.softmax(scores, dim=-1)
        # 将注意力权重与 value 进行加权求和
        attn_output = torch.bmm(attn_weights, value)
        
        # 将输出张量还原成原来的形状
        attn_output = attn_output.view(batch_size, self.num_heads, seq_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        return attn_output

以上代码中的 SelfAttention 类即为自注意力机制，其中 embed_dim 表示输入张量的维度，num_heads 表示要将输入张量分成多少份进行处理。在 forward 方法中，首先将输入张量 x 沿着最后一维分成 num_heads 份，然后进行线性变换（即计算 query、key、value），接着计算注意力权重，最后将注意力权重与 value 进行加权求和得到最终的输出张量。