用 PyTorch 实现 Transformer

项目结构

我们会实现 Encoder-Decoder 结构的 Transformer，包含：Embedding、Positional Encoding、Multi-Head Attention、Feed Forward 和完整的 Encoder/Decoder 层。

核心实现

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)

训练技巧

• 使用学习率预热
• 标签平滑防止过拟合
• 梯度累积应对显存不足

建议先阅读 Transformer 架构详解 理解理论基础，再对照代码实践。本文也是 从零构建 GPT 模型 的前置内容。