多头注意力机制 - 代码测试

import torch
from dm01_input import Embeddings, PositionalEncoding
from dm04_attention import MultiHeadAttention


def use_multihead_attn():
    """测试多头注意力机制"""
    # 1. 创建多头注意力层对象
    my_attention = MultiHeadAttention(512, 8)
    
    # 2. 获取位置编码处理后的结果
    position_x = use_position()  # (2, 4, 512)
    
    # 3. 定义 QKV（自注意力机制，Q=K=V）
    Q = K = V = position_x
    
    # 4. 创建掩码张量（可选）
    # 形状：(batch, 1, seq_len, seq_len)
    mask = torch.ones(2, 1, 4, 4)
    
    # 5. 送入模型获取结果
    result = my_attention(Q, K, V, mask)
    
    # 6. 打印输出形状
    print(f"多头注意力层结果形状: {result.shape}")
    
    # 7. 返回结果
    return result


def use_position():
    """获取位置编码处理后的结果"""
    VOCAB_SIZE = 1000
    D_MODEL = 512
    
    my_embedding = Embeddings(D_MODEL, VOCAB_SIZE)
    x = torch.LongTensor([[1, 2, 3, 4], [5, 6, 7, 8]])
    embedded_x = my_embedding(x)
    
    my_position = PositionalEncoding(D_MODEL, dropout=0.1, max_len=60)
    position_x = my_position(embedded_x)
    
    return position_x


# 测试
if __name__ == "__main__":
    result = use_multihead_attn()
    print(f"最终结果形状: {result.shape}")