Transformer 架构测试（上）

import copy

c = copy.deepcopy  # 不加小括号是赋值，加小括号是调用

# 使用
attn = MultiHeadAttention(d_model, num_heads)
self_attn = c(attn)  # 等价于 copy.deepcopy(attn)

import torch
import torch.nn as nn
import copy


class MultiHeadAttention(nn.Module):
    """多头注意力机制"""
    
    def __init__(self, d_model, num_heads, dropout=0.1):
        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)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, Q, K, V, mask=None):
        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)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        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, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        
        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):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)


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


def make_model(source_vocab, target_vocab, N=6, d_model=512, d_ff=2048, num_heads=8, dropout=0.1):
    """
    构建 Transformer 模型
    
    Args:
        source_vocab: 源词汇表大小
        target_vocab: 目标词汇表大小
        N: 编码器/解码器层数（默认 6）
        d_model: 词向量维度（默认 512）
        d_ff: 前馈全连接层维度（默认 2048）
        num_heads: 多头注意力头数（默认 8）
        dropout: 随机失活概率（默认 0.1）
    
    Returns:
        model: Transformer 模型
    """
    # 1. 定义深拷贝工具函数
    c = copy.deepcopy
    
    # 2. 创建词嵌入层
    source_embed = nn.Embedding(source_vocab, d_model)
    target_embed = nn.Embedding(target_vocab, d_model)
    
    # 3. 创建位置编码层
    position = PositionalEncoding(d_model, dropout)
    
    # 4. 创建多头注意力机制
    attn = MultiHeadAttention(d_model, num_heads, dropout)
    
    # 5. 创建前馈全连接层
    ff = FeedForward(d_model, d_ff, dropout)
    
    # 6. 创建编码器（使用 PyTorch 内置）
    encoder = nn.TransformerEncoder(
        nn.TransformerEncoderLayer(d_model, num_heads, d_ff, dropout, batch_first=True),
        N
    )
    
    # 7. 创建解码器（使用 PyTorch 内置）
    decoder = nn.TransformerDecoder(
        nn.TransformerDecoderLayer(d_model, num_heads, d_ff, dropout, batch_first=True),
        N
    )
    
    # 8. 创建源输入嵌入层（处理链：词嵌入 + 位置编码）
    source_embed_seq = nn.Sequential(c(source_embed), c(position))
    
    # 9. 创建目标输入嵌入层（处理链：词嵌入 + 位置编码）
    target_embed_seq = nn.Sequential(c(target_embed), c(position))
    
    # 10. 创建输出生成器
    generator = nn.Linear(d_model, target_vocab)
    
    # 11. 创建 Transformer 模型
    class EncoderDecoder(nn.Module):
        def __init__(self, encoder, decoder, source_embed, target_embed, generator):
            super(EncoderDecoder, self).__init__()
            self.encoder = encoder
            self.decoder = decoder
            self.source_embed = source_embed
            self.target_embed = target_embed
            self.generator = generator
        
        def forward(self, source_x, target_y, source_mask, target_mask):
            encoder_output = self.encode(source_x, source_mask)
            output = self.decode(target_y, encoder_output, source_mask, target_mask)
            return self.generator(output)
        
        def encode(self, source_x, source_mask):
            return self.encoder(self.source_embed(source_x), source_mask)
        
        def decode(self, target_y, encoder_output, source_mask, target_mask):
            return self.decoder(self.target_embed(target_y), encoder_output, source_mask, target_mask)
    
    model = EncoderDecoder(encoder, decoder, source_embed_seq, target_embed_seq, generator)
    
    # 12. 参数初始化（Xavier）
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    
    return model


# 测试
if __name__ == "__main__":
    print("=== Transformer 架构测试（上）===")
    
    source_vocab = 1000
    target_vocab = 1000
    
    # 构建 Transformer 模型
    model = make_model(source_vocab, target_vocab, N=6, d_model=512, num_heads=8)
    
    print(f"Transformer 模型创建成功")
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")