多头注意力机制 - 代码实现

def clones(module, n):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])

assert d_model % n_heads == 0  # 确保能整除

import torch
import torch.nn as nn
import math
import copy


def clones(module, n):
    """
    克隆 N 个相同模块
    :param module: 要被克隆的模块
    :param n: 克隆次数
    :return: 包含 N 个相同模块的列表（深拷贝）
    """
    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])


def attention(query, key, value, mask=None, dropout=None):
    """注意力计算"""
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = torch.softmax(scores, dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn


class MultiHeadAttention(nn.Module):
    """多头注意力机制"""
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        super().__init__()
        
        # 确保维度能整除
        assert d_model % n_heads == 0, f"d_model ({d_model}) 必须能被 n_heads ({n_heads}) 整除"
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每头维度
        
        # 4 个线性变换层
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        self.attn = None
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, Q, K, V, mask=None):
        """
        前向传播
        :param Q: 查询张量 (batch, seq_len, d_model)
        :param K: 键张量 (batch, seq_len, d_model)
        :param V: 值张量 (batch, seq_len, d_model)
        :param mask: 掩码张量
        :return: 多头注意力输出
        """
        batch_size = Q.size(0)
        
        # 1. 线性变换（投影）
        Q = self.W_q(Q)  # (batch, seq_len, d_model)
        K = self.W_k(K)
        V = self.W_v(V)
        
        # 2. 分头：(batch, seq_len, d_model) -> (batch, n_heads, seq_len, d_k)
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # 3. 计算注意力（每个头并行计算）
        output, self.attn = attention(Q, K, V, mask=mask, dropout=self.dropout)
        
        # 4. Concat：(batch, n_heads, seq_len, d_k) -> (batch, seq_len, d_model)
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, -1, self.d_model)
        
        # 5. 最终线性变换
        return self.W_o(output)


# 测试代码
if __name__ == "__main__":
    # 测试 clones 函数
    linear = nn.Linear(512, 512)
    cloned = clones(linear, 6)
    print(f"克隆数量: {len(cloned)}")
    
    # 测试 MultiHeadAttention
    mha = MultiHeadAttention(d_model=512, n_heads=8, dropout=0.1)
    Q = K = V = torch.randn(2, 4, 512)
    output = mha(Q, K, V)
    
    print(f"输入形状: {Q.shape}")
    print(f"输出形状: {output.shape}")
    print(f"d_model: {mha.d_model}")
    print(f"n_heads: {mha.n_heads}")
    print(f"d_k: {mha.d_k}")