代码题：Multi-Head Attention 实现

import math
import torch
from torch import nn


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout_rate):
        super().__init__()
        self.d_model = d_model          # 模型总维度
        self.n_heads = n_heads          # 注意力头数
        self.d_head = d_model // n_heads  # 每个头的维度

        # 四个线性变换层
        self.q_linear = nn.Linear(d_model, d_model)   # Q 的投影矩阵
        self.k_linear = nn.Linear(d_model, d_model)   # K 的投影矩阵
        self.v_linear = nn.Linear(d_model, d_model)   # V 的投影矩阵
        self.out_linear = nn.Linear(d_model, d_model) # 输出投影矩阵

        self.scale = math.sqrt(self.d_head)           # 缩放因子
        self.dropout = nn.Dropout(dropout_rate)       # Dropout 层

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # 1. 线性变换 + 多头拆分
        # Q: (batch, seq_len, d_model) -> (batch, n_heads, seq_len, d_head)
        q = (
            self.q_linear(query)                           # (batch, seq_len, d_model)
            .view(batch_size, -1, self.n_heads, self.d_head) # (batch, seq_len, n_heads, d_head)
            .transpose(1, 2)                                # (batch, n_heads, seq_len, d_head)
        )
        k = (
            self.k_linear(key)
            .view(batch_size, -1, self.n_heads, self.d_head)
            .transpose(1, 2)
        )
        v = (
            self.v_linear(value)
            .view(batch_size, -1, self.n_heads, self.d_head)
            .transpose(1, 2)
        )

        # 2. 计算缩放点积注意力分数
        atten_scores = (q @ k.transpose(-1, -2)) / self.scale  # (batch, n_heads, q_len, k_len)

        # 3. Mask 处理（将 padding 位置设为 -inf）
        if mask is not None:
            if mask.dim() == 2:                     # 如果 mask 维度为 (seq_len, seq_len)
                mask = mask.unsqueeze(0).unsqueeze(0)  # 扩展为 (1, 1, seq_len, seq_len)
            atten_scores = atten_scores.masked_fill(mask, float("-inf"))

        # 4. Softmax + Dropout
        atten_weights = atten_scores.softmax(dim=-1)      # (batch, n_heads, q_len, k_len)
        droped_atten_weights = self.dropout(atten_weights)

        # 5. 加权求和：权重 × V
        atten_output = droped_atten_weights @ v           # (batch, n_heads, q_len, d_head)

        # 6. 合并多头：transpose + view
        multi_atten_output = atten_output.transpose(1, 2).view(
            batch_size, -1, self.d_model
        )                                                 # (batch, q_len, d_model)

        # 7. 输出线性变换
        output = self.out_linear(multi_atten_output)
        return output, atten_weights

if __name__ == "__main__":
    # 初始化参数
    batch_size = 2
    seq_len = 4
    d_model = 8
    n_heads = 2
    dropout_rate = 0.1

    mha = MultiHeadAttention(d_model, n_heads, dropout_rate)

    # 生成随机输入
    query = torch.randn(batch_size, seq_len, d_model)
    key = torch.randn(batch_size, seq_len, d_model)
    value = torch.randn(batch_size, seq_len, d_model)

    # 无 mask 测试
    output, attn = mha(query, key, value)
    print(f"无 mask 输出形状: {output.shape}")  # 期望: (2, 4, 8)
    print(f"注意力权重形状: {attn.shape}")       # 期望: (2, 2, 4, 4)

    # 带 mask 测试（模拟 padding mask）
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()  # 上三角 mask
    output_masked, attn_masked = mha(query, key, value, mask=mask)
    print(f"带 mask 输出形状: {output_masked.shape}")  # 期望: (2, 4, 8)
    print(f"带 mask 注意力权重:\n{attn_masked[0, 0]}") # 查看第一头注意力（上三角应为0）