摆摊卖爱情我很早就想写一篇“用 PyTorch 从零实现一个 Transformer”的经验总结，原因很简单：几乎每个入门深度学习的人都会卡在“看得懂图、写不出代码”这一步。Transformer 的架构图在网上随处可见，Encoder、Decoder、Multi-Head Attention 这些名词大家都能脱口而出，可真要关掉教程，自己从空目录开始写一个能训练、能推理的模型，很多人会在一开始就懵掉——位置编码怎么加？Mask 的形状到底是什么？Decoder 的输入往右移一位是什么意思？这些细节，看文章的时候都觉得“懂了”，动手才会发现全是坑。
这篇文章我尽量按实际项目的推进顺序来写，记录我从任务设计、数据处理到核心模块实现、训练调试的完整过程。代码全部只用 PyTorch 的基础组件，不碰 nn.Transformer、nn.MultiheadAttention 这些高层封装。适合两种人看：一种是已经看过 Transformer 图解、想动手验证理解的初学者，另一种是面试前需要快速把整个模型串起来的老手。我会把每个关键选择背后的原因也一并讲清楚，这样你不仅能照抄代码，还能知道每一步到底在做什么。
1. 为什么非得“手撕”一遍 Transformer
1.1 纸上得来终觉浅，注意力机制尤其如此
Transformer 的核心是自注意力机制，这个机制本身的公式很简单：Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V。但我观察到一个现象：很多人能把公式默写出来，却回答不了几个特别基础的问题——为什么除的是 sqrt(d_k) 而不是 d_k？为什么 Q、K、V 要拆成多个头再拼回去？训练的时候 padding 位置和未来位置是怎么被“屏蔽”掉的？这些问题靠看是看不会的，只有自己实现一遍，才会在写错、调试、看 loss 变化的过程中真正留下肌肉记忆。
我这次实现还有一个原则：能自己写的地方绝不调现成 API。nn.MultiheadAttention 确实方便，但正因为方便，它把很多细节掩盖掉了。比如输入张量是 [batch, seq_len, embed_dim]，进到这个模块内部后会被 reshape 成 [seq_len, batch, embed_dim]，这种维度变化如果不自己处理一次，很难对 PyTorch 中 attention 的 tensor 流转建立直观感受。手写虽然代码量多一点，但换来的是对整个数据流清晰到每个维度的掌控感。
1.2 直接调用 nn.Transformer 差在哪
PyTorch 官方提供了 nn.Transformer，它可以一行代码搭出 Transformer，但这恰恰是问题所在。你用它完成一个简单任务后，除了会填参数，很难说真正理解了这个模型。参数填错了，报错信息又长又绕，你根本不知道问题出在 Encoder 还是 Decoder，更别提定位到某个维度 mismatch 的具体原因。
从我带过项目的经验来看，能徒手写出 Transformer 的人，遇到模型不收敛、loss 变成 NaN、推理结果全是一个 token 这类问题时，定位速度会快得多。因为他脑子里有一张完整的计算图，知道数据在每一层之后长什么样。这就好比老司机能通过发动机声音判断大致故障，而只会踩油门的人只能原地等救援。
1.3 这个项目适合谁、最终能带走什么
这个项目从零开始包括任务设计、数据构造、模型实现、训练评估，整套流程走下来大约半天到一天时间。不需要 GPU，CPU 上几分钟就能完成训练收敛，关键是能跑通、能验证、能观察注意力权重的变化。如果你是第一次接触 Transformer，建议把代码从头到尾敲一遍，不要直接复制粘贴。敲代码这个动作本身，就是逼迫大脑处理每个细节的过程。
需要模型API调用？ 免费领10W Token，多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 任务选型：用“序列累加”验证模型
2.1 为什么选累加而不是翻译或文本生成
从零实现一个模型，最怕的是选错验证任务。一开始我考虑过用英文到中文的翻译，但翻译任务对数据量、词表、训练时间的要求太高，在 CPU 上等一次完整训练可能要按天计算，非常不利于快速迭代调 bug。后来我换成了语言模型式的人物对话生成，又发现评估比较主观，模型输出一句不通顺的话你很难判断是模型没训练好，还是实现本身有 bug。
最后我选择了一个特别简单的合成任务：给定一串数字，模型输出每个位置之前的累加和。举个例子，输入 [3, 1, 4, 1, 5]，输出 [3, 4, 8, 9, 14]。这个任务有几个非常适合验证模型实现正确性的特点：
- 输入输出都是离散 token，天然适配分类式的交叉熵损失。
- 输出长度和输入长度完全相同，不需要复杂的长度控制逻辑。
- 任务本身要求模型“记住”前缀和，这意味着 Decoder 必须在每个时间步都 attend 到前序所有位置，能有效检验 attention 机制是否真的在起作用。
- 数据可以完全合成，批量生成只需要几行代码，不用折腾下载数据集。
- 训练速度快，CPU 上几分钟就能看到收敛。
如果你是一个有经验的工程师，可能会觉得这个任务“太简单了，一点都不像真实项目”。但请记住，我们的目标不是刷榜，而是用最短路径验证手写 Transformer 的正确性。模型能在这个任务上收敛到接近 100% 准确率，说明位置编码、多头注意力、mask 这些模块大概率没错；如果这些模块有问题，再复杂的任务也无法成功。
2.2 输入输出与词表设计
词表设计要同时考虑输入和输出。输入是 0~9 的数字，但输出是累加和，范围会超过 9。比如最长序列长度我设为 12，输入最大为 9，则最大的累加和是 12 × 9 = 108。所以输出词表至少要覆盖 0 到 108 这些数字。
为了简单，我不区分输入词表和输出词表，统一用一个包含 0~127 的数字 token 空间，再加上几个特殊 token 就够了。具体分配如下：
| token 含义 | 索引范围 |
|---|---|
| 数字 0~127 | 0~127 |
| BOS（begin of sequence） | 128 |
| PAD | 129 |
最大累加和是 108，低于 127，因此用这个词表不会出现越界。vocab_size = 130。输入数字只会用到 0~9 这部分，但词表留出更大的空间可以让输出 token 的范围覆盖所有可能的累加结果，实际训练中也更省心，省去了判断输出到底会不会越界的麻烦。
2.3 模型配置与参数量估算
模型规模不能太大，否则 CPU 上训练会让人失去耐心；也不能太小，否则 attention 机制发挥不出效果。我选择的配置如下：
| 超参数 | 取值 |
|---|---|
| d_model（embedding 维度） | 64 |
| num_heads（注意力头数） | 4 |
| d_ff（前馈层中间维度） | 128 |
| num_layers（编码器和解码器层数） | 2 |
| dropout | 0.1 |
| max_len（位置编码最大长度） | 64 |
| vocab_size | 130 |
关于头数的选择，我特意选了 4 而不是常见的 8。因为 d_model 是 64，如果分成 8 个头，每个头的维度只剩 8，信息容量有点紧张；4 个头时每个头有 16 维，相对合理。这个配置下，模型的参数量大概在 30 万左右，非常轻量。训练 30 个 epoch，CPU 上只需要两三分钟。
3. 数据准备：为 Transformer 造一份合适的数据集
3.1 合成数据的生成逻辑
数据准备是很多人容易轻视的环节，但它对训练效果的影响极其直接。我用 PyTorch 的 Dataset 类来封装数据生成逻辑，核心思路是：每次随机生成长度在 4 到 12 之间的数字序列，每个数字取值 0 到 9，然后计算前缀和作为目标序列。
python复制import random
import torch
from torch.utils.data import Dataset
PAD_TOKEN = 129
BOS_TOKEN = 128
MAX_VAL = 127
class CumSumDataset(Dataset):
def __init__(self, num_samples=10000, min_len=4, max_len=12, seed=42):
super().__init__()
self.num_samples = num_samples
self.min_len = min_len
self.max_len = max_len
random.seed(seed)
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
length = random.randint(self.min_len, self.max_len)
src = [random.randint(0, 9) for _ in range(length)]
# 前缀和即为目标输出
tgt = []
s = 0
for x in src:
s += x
tgt.append(s)
re