PyTorch 多头注意力机制实现详解：代码、注释和维度分析

下面是一个简单的 PyTorch 多头注意力机制的实现，包含了注释和变量维度。

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, input_dim, key_dim, value_dim):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.key_dim = key_dim
        self.value_dim = value_dim
        
        # 定义线性层，用于将输入数据转换为 Q、K、V
        self.linear_q = nn.Linear(input_dim, num_heads * key_dim)
        self.linear_k = nn.Linear(input_dim, num_heads * key_dim)
        self.linear_v = nn.Linear(input_dim, num_heads * value_dim)
        
        # 定义线性层，用于将多头注意力的输出合并
        self.linear_out = nn.Linear(num_heads * value_dim, input_dim)
        
    def forward(self, x):
        batch_size, seq_len, input_dim = x.shape
        
        # 将输入数据转换为 Q、K、V
        q = self.linear_q(x).view(batch_size, seq_len, self.num_heads, self.key_dim)
        k = self.linear_k(x).view(batch_size, seq_len, self.num_heads, self.key_dim)
        v = self.linear_v(x).view(batch_size, seq_len, self.num_heads, self.value_dim)
        
        # 计算注意力分数
        score = torch.einsum('bqhd,bkhd->bhqk', q, k) / torch.sqrt(torch.tensor(self.key_dim, dtype=torch.float32))
        # score 的维度为 [batch_size, num_heads, seq_len, seq_len]
        
        # 计算注意力权重
        attention_weights = torch.softmax(score, dim=-1)
        # attention_weights 的维度为 [batch_size, num_heads, seq_len, seq_len]
        
        # 计算注意力输出
        attention_output = torch.einsum('bhqk,bkhd->bqhd', attention_weights, v)
        # attention_output 的维度为 [batch_size, seq_len, num_heads, value_dim]
        attention_output = attention_output.view(batch_size, seq_len, -1)
        # attention_output 的维度为 [batch_size, seq_len, num_heads * value_dim]
        
        # 合并多头注意力的输出
        output = self.linear_out(attention_output)
        # output 的维度为 [batch_size, seq_len, input_dim]
        
        return output

下面是每个变量的维度：

x: 输入数据，维度为 [batch_size, seq_len, input_dim]
num_heads: 多头注意力的头数
input_dim: 输入数据的维度
key_dim: Q、K 的维度（也就是注意力分数的维度）
value_dim: V 的维度（也就是注意力输出的维度）
linear_q: 将输入数据转换为 Q 的线性层，输出维度为 [batch_size, seq_len, num_heads * key_dim]
linear_k: 将输入数据转换为 K 的线性层，输出维度为 [batch_size, seq_len, num_heads * key_dim]
linear_v: 将输入数据转换为 V 的线性层，输出维度为 [batch_size, seq_len, num_heads * value_dim]
score: 注意力分数，维度为 [batch_size, num_heads, seq_len, seq_len]
attention_weights: 注意力权重，维度为 [batch_size, num_heads, seq_len, seq_len]
attention_output: 注意力输出，维度为 [batch_size, seq_len, num_heads, value_dim]
output: 多头注意力的输出，维度为 [batch_size, seq_len, input_dim]