解码器层 - 代码测试

import copy

# 多头注意力层
multihead_attn = MultiHeadAttention(...)
self_attn = copy.deepcopy(multihead_attn)  # 自注意力
src_attn = copy.deepcopy(multihead_attn)  # 编码器-解码器注意力

import torch
import torch.nn as nn
import copy


class MultiHeadAttention(nn.Module):
    """多头注意力机制"""
    
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        
        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.fc = nn.Linear(d_model, d_model)
    
    def forward(self, Q, K, V):
        batch_size = Q.size(0)
        
        Q = self.w_q(Q).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.w_k(K).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.w_v(V).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)
        
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.fc(output)
        
        return output


class PositionalEncoding(nn.Module):
    """位置编码层"""
    
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]


class FeedForward(nn.Module):
    """前馈全连接层"""
    
    def __init__(self, d_model, d_ff=2048):
        super(FeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
    
    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))


def use_decoder_layer():
    """测试解码器层"""
    print("=== 解码器层 - 代码测试 ===")
    
    # 1. 定义变量记录解码器输入（Q）
    vocab_size = 4345
    d_model = 512
    num_heads = 8
    batch_size = 1
    seq_len = 10
    
    Y = torch.randint(0, vocab_size, (batch_size, seq_len))  # [1, 10]
    print(f"输入 Y 形状: {Y.shape}")
    
    # 2. 创建词嵌入层
    embedding = nn.Embedding(vocab_size, d_model)
    
    # 3. 将 Y 输入变成词向量形式
    embedded = embedding(Y)  # [1, 10, 512]
    print(f"词嵌入后形状: {embedded.shape}")
    
    # 4. 创建位置编码层，处理 Y
    my_position = PositionalEncoding(d_model)
    position_y = my_position(embedded)  # [1, 10, 512]
    print(f"位置编码后形状: {position_y.shape}")
    
    # 5. 创建多头注意力层（使用深拷贝）
    multihead_attn = MultiHeadAttention(d_model, num_heads)
    self_attn = copy.deepcopy(multihead_attn)  # 自注意力
    src_attn = copy.deepcopy(multihead_attn)  # 编码器-解码器注意力
    
    # 6. 创建前馈全连接层
    ff = FeedForward(d_model)
    
    # 7. 构建三个子层
    print(f"\n三个子层:")
    print(f"  子层 1（掩码多头自注意力）: {self_attn}")
    print(f"  子层 2（编码器-解码器注意力）: {src_attn}")
    print(f"  子层 3（前馈全连接层）: {ff}")
    
    # 8. 测试
    # 子层 1：掩码多头自注意力
    output1 = self_attn(position_y, position_y, position_y)
    print(f"\n子层 1 输出形状: {output1.shape}")  # [1, 10, 512]
    
    # 子层 2：编码器-解码器注意力
    encoder_output = torch.randn(batch_size, seq_len, d_model)  # 编码器输出
    output2 = src_attn(output1, encoder_output, encoder_output)
    print(f"子层 2 输出形状: {output2.shape}")  # [1, 10, 512]
    
    # 子层 3：前馈全连接层
    output3 = ff(output2)
    print(f"子层 3 输出形状: {output3.shape}")  # [1, 10, 512]


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