自注意力机制 pytorch代码

自注意力机制（Self-Attention Mechanism），也被称为Transformer的基石，是一种在处理序列数据时非常关键的技术，特别是在自然语言处理领域。它允许模型对输入序列中的每个位置都与其他位置的信息进行交互。在PyTorch中，我们可以使用`nn.Module`类来实现这个机制。

下面是一个简单的自注意力模块（Multi-Head Attention）的示例代码：

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.fc_q = nn.Linear(d_model, d_model) # Query Linear layer
        self.fc_k = nn.Linear(d_model, d_model) # Key Linear layer
        self.fc_v = nn.Linear(d_model, d_model) # Value Linear layer
        self.fc_out = nn.Linear(d_model, d_model)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Linear projections
        query = self.fc_q(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        key = self.fc_k(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        value = self.fc_v(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot product attention
        scores = torch.matmul(query, key.transpose(-2, -1)) / (self.head_dim ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -float('inf'))
        attn_weights = nn.functional.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, value)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        
        # Final linear projection
        output = self.fc_out(output)
        
        return output, attn_weights

在这个例子中，你需要提供查询（query）、键（key）和值（value）作为输入，其中query、key和value通常来自前一层的隐藏状态。`mask`用于防止当前位置看到后面的词。