从零实现教学级Transformer：PyTorch核心代码详解

1. 项目概述

今天我想分享一个教学级的Transformer代码实现，这是我为深度学习课程准备的实践案例。Transformer架构自2017年提出以来，已经成为NLP领域的基石模型，但很多初学者对其实现细节仍感到困惑。这个实现从零开始构建，去除了工业级框架的复杂性，保留了核心架构的完整性。

提示：这个实现使用Python和PyTorch，适合有一定深度学习基础的学习者。完整代码约800行，包含详细的注释和中间变量检查点。

2. 核心架构解析

2.1 自注意力机制实现

自注意力是Transformer最核心的组件。我们的实现采用标量计算方式展示计算过程：

python复制def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: 查询矩阵 [batch_size, num_heads, seq_len, d_k]
    K: 键矩阵
    V: 值矩阵
    """
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    attention = torch.softmax(scores, dim=-1)
    return torch.matmul(attention, V)

关键细节：

1. 缩放因子1/√d_k防止梯度消失
2. 使用mask处理变长序列
3. 多头注意力的并行计算实现

2.2 位置编码方案

Transformer没有RNN的时序结构，需要显式的位置编码：

python复制class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(1)]

注意：位置编码的维度必须与词嵌入维度一致，且不需要训练。

3. 完整模型搭建

3.1 编码器实现

编码器由N个相同层堆叠而成，每层包含：

1. 多头自注意力子层
2. 前馈网络子层
3. 残差连接和层归一化

python复制class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask):
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))
        ff_output = self.feed_forward(x)
        return self.norm2(x + self.dropout(ff_output))

3.2 解码器特殊处理

解码器相比编码器增加了：

1. 带掩码的多头自注意力（防止信息泄露）
2. 编码器-解码器注意力层

python复制class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, encoder_output, src_mask, tgt_mask):
        # 自注意力（带掩码）
        attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        # 编码器-解码器注意力
        attn_output = self.cross_attn(x, encoder_output, encoder_output, src_mask)
        x = self.norm2(x + self.dropout(attn_output))
        
        ff_output = self.feed_forward(x)
        return self.norm3(x + self.dropout(ff_output))

4. 训练技巧与调试

4.1 学习率调度器

Transformer使用特殊的学习率预热策略：

python复制class WarmupScheduler:
    def __init__(self, optimizer, d_model, warmup_steps=4000):
        self.optimizer = optimizer
        self.d_model = d_model
        self.warmup_steps = warmup_steps
        self.current_step = 0
    
    def step(self):
        self.current_step += 1
        lr = (self.d_model ** -0.5) * min(
            self.current_step ** -0.5,
            self.current_step * self.warmup_steps ** -1.5
        )
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

4.2 常见问题排查

1. 梯度爆炸：

   • 检查注意力分数缩放
   • 添加梯度裁剪
   • 适当减小初始化范围

2. 模型不收敛：

   • 验证输入输出维度匹配
   • 检查mask是否正确应用
   • 确认学习率预热生效

3. 显存不足：

   • 减小batch_size
   • 使用梯度累积
   • 混合精度训练

5. 扩展应用方向

这个基础实现可以扩展为：

1. 视觉Transformer：

   • 将图像分块作为输入序列
   • 添加CNN特征提取层

2. 时间序列预测：

   • 修改位置编码适应时序数据
   • 添加季节性特征嵌入

3. 多模态模型：

   • 文本和图像共享注意力机制
   • 设计跨模态融合层

我在实际教学中发现，通过这个简化实现，学生能更清晰地理解：

• 注意力权重的动态计算过程
• 位置编码的时空特性
• 残差连接对深层网络的稳定作用