注意力机制 - 测试代码

import torch
import torch.nn as nn
import torch.nn.functional as F


class MyAttention(nn.Module):
    """自定义注意力机制模块"""
    
    def __init__(self, query_size, key_size, value_size, output_size):
        super(MyAttention, self).__init__()
        self.query_layer = nn.Linear(query_size, key_size)
        self.attn_weights = None
        self.output_layer = nn.Linear(value_size, output_size)
    
    def forward(self, Q, K, V):
        # 1. 计算注意力权重分布
        Q = self.query_layer(Q)
        attn_weights = F.softmax(torch.matmul(Q, K.transpose(-2, -1)), dim=-1)
        self.attn_weights = attn_weights
        
        # 2. 用 bmm() 计算注意力结果表示
        attention_q = torch.bmm(attn_weights, V)
        
        # 3. 按指定维度输出
        output = self.output_layer(attention_q)
        
        return output, attn_weights


def test_attention():
    """测试注意力机制"""
    print("=== 注意力机制 - 测试代码 ===")
    
    # 1. 实例化参数设置
    query_size = 32
    key_size = 32
    value_size = 64
    output_size = 32
    
    # 2. 创建随机输入张量
    Q = torch.randn(1, 1, query_size)  # 查询张量 [批次, 序列, 特征]
    K = torch.randn(1, 32, key_size)  # 键张量 [批次, 单词数, 特征]
    V = torch.randn(1, 32, value_size)  # 值张量 [批次, 单词数, 词向量维度]
    
    print(f"Q 形状: {Q.shape}")
    print(f"K 形状: {K.shape}")
    print(f"V 形状: {V.shape}")
    
    # 3. 实例化注意力机制模块并测试
    my_attn = MyAttention(query_size, key_size, value_size, output_size)
    output, attn_weights = my_attn(Q, K, V)
    
    # 4. 输出结果
    print("\n" + "=" * 50)
    print(f"查询张量 Q 的注意力结果表示:")
    print(f"  output 形状: {output.shape}")  # [1, 1, 32]
    print(f"  output 内容: {output}")
    
    print(f"\n查询张量 Q 的注意力权重分布:")
    print(f"  attn_weights 形状: {attn_weights.shape}")  # [1, 1, 32]
    print(f"  attn_weights 内容: {attn_weights}")
    print(f"  权重和: {attn_weights.sum().item():.4f}")  # 应该接近 1


# 测试
if __name__ == "__main__":
    test_attention()