1. 项目思路拆解：为什么用PyTorch + Transformer做入门
先说结论：如果你想在PyTorch里跑通Transformer，又不想一上来就啃论文、看几十个前置概念，这个小项目是很好的起点。我用一个几十行的模型，不依赖任何预训练权重和现成翻译库，让模型自己学会了把一个整数序列倒序输出。任务本身很弱智，但它能把Transformer最核心的机制——注意力、多头、位置编码、mask——全部串起来，跑通之后你会对编码器-解码器结构有非常具体的体感，而不是停留在“听过、见过、不会用”的层面。
1.1 选型逻辑：为什么偏偏是这两个组合
PyTorch和Transformer几乎是现在做序列模型最省心的组合。PyTorch的动态图机制让调试变得非常直观，你可以在forward中间随意打印张量形状、检查数值分布，这一点在理解和调Transformer这种结构复杂的模型时尤其重要。很多初学者卡在维度对不上、mask不知道怎么传这类问题上，如果用静态图框架，光排查这些就要绕很多弯路，但PyTorch里你可以直接打印每一层的输入输出，问题一眼就能定位。
再说Transformer本身。这两年做序列建模，LSTM和GRU逐渐退出主流，核心原因不是精度问题，而是Transformer有三个LSTM比不了的优势：第一，训练可以完全并行，LSTM必须逐步展开，GPU利用率天然受限；第二，注意力机制直接建模任意两个位置的关系，长距离依赖的捕捉能力比循环结构强一大截；第三，它是一切大模型的基础架构，理解了小Transformer，你再看GPT、BERT这类模型，至少骨架上是熟悉的。
1.2 示例任务设计：怎么用最小成本验证模型真的学会了
这个示例里我选用的是“序列倒序反转任务”：给模型输入一串随机整数，比如[3, 5, 1, 4, 2]，要求模型输出[2, 4, 1, 5, 3]。为什么选这个任务？因为它有三个好处。
第一，数据生成不需要任何外部资源。自己写一个随机数生成函数就行，不需要下载数据集、不需要做清洗和预处理，几行代码就能造出无限训练样本。第二，倒序任务对模型的序列建模能力要求其实很高——模型必须记住输入序列的每一个位置，并且把它们的顺序关系完全反过来，这比简单的复制任务难很多，早期RNN在这个任务上训练非常慢，而Transformer能很好地处理。第三，它能直观检验模型是否真的“理解”了序列结构，如果模型只是死记硬背，那换个长度的序列就会原形毕露。
1.3 预期效果与目标边界
跑通这个示例后，你应该达到三个目标：一是能手写一个完整的编码器-解码器Transformer结构，而不是只会调用现成的nn.Transformer然后黑盒使用；二是理解位置编码为什么存在、mask到底在干什么；三是遇到训练不收敛或者loss变成NaN时，知道从哪些方向排查。
需要明确的是，这个示例不做大规模语料、不做束搜索、不做预训练。它就是一台“解剖台上的麻雀”，让你看清楚Transformer是怎么工作的。等你想真正做机器翻译或文本生成时，再在这个基础上替换数据、加优化技巧，是完全平滑的过渡。
2. 核心机制拆解：Transformer到底在做什么
这一节我会用最通俗的方式把Transformer的核心组件讲清楚。很多人学Transformer最大的问题是：每个概念单独看都懂，但不知道它们是怎么配合的。所以我换个讲法，用一个日常场景把整条数据流串起来。
2.1 注意力机制：从“查字典”到“找关联”
注意力机制的中心思想，可以用一句话概括：对序列里的每一个位置，我都要和所有其他位置做一次“关联度打分”，然后按照关联度把信息汇聚起来。生活里最贴切的类比是查字典——你要理解某个词在句子里是什么意思，不是只看这个词本身，而是要结合它周围的词来判断。比如“苹果”这个词，如果句子里有“华为”“手机”，它指的是数码产品；如果有“红富士”“削皮”，它指的就是水果。Transformer做的是同一件事，只不过它更彻底，它让序列里任意两个位置都互相“看”一眼。
具体实现上，每个位置的token会生成三个向量：Query（查询）、Key（索引）、Value（内容）。打分过程就是把Query和所有位置的Key做点积，得到的分数经过softmax归一化，变成一组权重，然后用这组权重去加权求和所有位置的Value。这里有个容易忽略的细节：点积的结果要除以一个缩放因子（根号下d_k），原因是在维度较高时点积的数值会变得很大，softmax会进入饱和区，梯度就消失了。这个缩放因子既不是经验值也不是可学习参数，背后是有严格的数学推导的，你只需要记住它。
2.2 多头注意力：为什么“多头”比“单头”强
多头注意力就是把拆成多个子空间分别做注意力。比如d_model=128、nhead=4，那就把128维切成4份，每份32维，独立做一次注意力，最后把4份结果拼回128维。
为什么要这么麻烦？一段话里往往同时存在多组关联关系，比如翻译句子时，一个词既想关注主语是谁，又想关注时态发生了什么，还想关注后面接什么宾语。单头注意力只能学习一种关系权重，容易顾此失彼。多头相当于让模型同时从多个不同角度观察输入，每个头负责捕捉一种关联模式。有人说“每个头学到了不同的语法或语义特征”，这个说法不完全准确，但直觉上是对的——不同的头确实会倾向于关注不同类型的依赖关系。
实操里还有一个容易踩的坑：在PyTorch中自己实现多头注意力时，很多人会把head和batch维度搞混。正确思路是先做线性变换得到Q、K、V，然后reshape成(batch, seq_len, nhead, head_dim)，再调换维度变成(batch, nhead, seq_len, head_dim)，最后用torch.matmul做批量矩阵乘法。如果你选择直接使用nn.MultiheadAttention，记得注意batch_first这个参数——PyTorch 1.6之后的版本默认batch_first=False，也就是说输入是(seq_len, batch, embed_dim)，不设对的话形状很容易糊里糊涂报错。
2.3 位置编码：给Transformer补上“顺序感”
Transformer最大的缺陷是没有天然的序列顺序概念。RNN按时间步逐步输入，天然知道谁先谁后，但Attention的计算是两个位置之间的直接交互，如果把同一个序列的所有token顺序打乱，算出来的结果完全一样，这显然不对。语言是高度依赖语序的，“我打你”和“你打我”意思完全不同。
解决方法是往输入里注入位置信息，也就是位置编码。这里用的是一个固定公式：每个位置上，偶数维度用正弦函数、奇数维度用余弦函数，频率从低到高逐维递增。这样设计的一个好处是，位置之间的距离可以通过内积反映出来，而且任意位置都能由一组周期性函数唯一表示，不需要额外学习参数。
实际操作中，nn.Transformer内部不会自动帮你加位置编码，你必须自己实现并加在Embedding后面。要注意两个细节：一是位置编码矩阵要用register_buffer而不是作为普通参数注册，这样它不会随梯度更新，但能跟着模型一起迁移到GPU上；二是要对max_len设置上限，推理时如果输入长度超过训练时的max_len，就会发生切片越界，报错会很隐蔽——它往往不是直接说“越界”，而是维度不匹配。
2.4 Mask机制：三种mask各管哪一摊
Mask是Transformer里最容易糊弄过去的部分，但也是影响成败的关键。常用的mask有三种，它们的用途完全不同。
第一种是padding mask。一个batch里不同样本长度不同，短样本要补padding符号到统一长度，但注意力不能把这些padding位置算进去，否则等于告诉模型“这里有无效信息，但你可以关注它”。这种mask的shape是(batch, seq_len)，值为True的位置表示该位置是padding，在计算softmax之前要把这些位置的分数设成负无穷。
第二种是目标序列mask，也叫防止未来信息泄露的mask。解码器生成时是一个词一个词往后生成的，第t个位置不应该看到第t+1个位置的内容。这个mask的shape是(seq_len, seq_len)，右上角为True，保证每个位置只能attend到它自己和它左边的位置。忘记加这个mask的后果非常隐蔽：训练时loss会降得飞快，准确率看起来很高，但一到推理阶段生成的内容就开始胡言乱语。因为训练时模型“偷看”了后面位置的答案，学到了一个作弊函数。
第三种是编码器-解码器交叉注意力的一些辅助mask，通常复用padding mask，让解码器只关注编码器输出的有效非padding位置。很多人上来就把三种mask混在一起，后面排查时非常痛苦。我后面会专门写一个排查表，告诉你每个mask的形状和用途，这个模块跳过的坑太多了。
3. 实操过程：从零到一搭建一个小型Transformer
这节是全文的核心。我会给出一个完整、可复现的实验过程，包括数据生成、模型定义、训练循环和推理验证。所有代码我都实际跑过，用的是PyTorch 2.x，CPU上几十秒就能看到效果，不需要GPU。
3.1 数据准备：用随机序列生成无限训练样本
我们的倒序任务核心函数非常简单，但需要提前定义三个特殊符号：PAD=0代表padding符，SOS=1代表起始符，EOS=2代表结束符。有效数字从3开始，这样能避免和特殊符号冲突。
import torch import torch.nn as nn import math import random PAD = 0 SOS = 1 EOS = 2 def generate_batch(batch_size=32, vocab_size=20, max_len=8): # 有效token范围是[3, vocab_size-1] src = torch.randint(3, vocab_size, (batch_size, max_len)) # 目标序列 = 源序列的倒序 rev_src = torch.flip(src, dims=[1]) # 解码器输入：前面加SOS tgt_input = torch.cat([torch.full((batch_size, 1), SOS, dtype=torch.long), rev_src], dim=1) # 学习目标：后面加EOS tgt_output = torch.cat([rev_src, torch.full((batch_size, 1), EOS, dtype=torch.long)], dim=1) return src, tgt_input, tgt_output这里有一点需要仔细理解：训练解码器时用的是一种叫“teacher forcing”的策略。每一步解码的输入不是上一步的预测结果，而是真实的目标序列（shifted right）。shifted right指的是目标序列向右平移一位：预测位置i时，输入是位置i之前的真实目标token。这就是为什么要在倒序序列前加一个SOS。比如真实目标是[4, 1, 5, 3]，那解码器输入就是[SOS, 4, 1, 5]，输出是[4, 1, 5, 3]对应的预测，最后一个位置预测EOS。整个过程逻辑要对齐——输出长度和输入长度一样，都等于max_len+1。
3.2 位置编码与模型主体定义
位置编码的代码非常标准，不断有初学者纠结为什么要用torch.arange(0, d_model, 2)和torch.arange(0, d_model, 2)这种写法。原因是要对偶数维和奇数维分别计算sin和cos，所以步长取2，每次跳过一个维度。div_term是频率系数，-math.log(10000.0) / d_model这个式子让频率随维度增大而指数衰减，最终得到的高维维度对应低频信号，低维维度对应高频信号，和信号处理中的傅里叶特征有相似之处。
class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=100): super().__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() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) # 注册为buffer，不参与梯度更新但会随模型迁移设备 self.register_buffer('pe', pe) def forward(self, x): return x + self.pe[:, :x.size(1)]模型主体我选择直接使用nn.Transformer这个封装类，但自己接上前面的Embedding、位置编码和最后的输出层。这样做的优点是结构清晰，代码量少，同时保留了手动干预的灵活性。有几个关键参数需要仔细说明。
d_model=64是Embedding的维度，这个值不能太小，否则表示能力不够；也不能太大，否则CPU上训练太慢。nhead=4是注意力头数，注意d_model必须能被nhead整除，这是很多人一开始报错的原因——维度除不断。num_encoder_layers=2、num_decoder_layers=2都取2，这个任务本身简单，没必要堆层数。dim_feedforward=256是前馈网络的隐藏层维度，Transformer每个注意力层后面会跟一个两层全连接网络，常见设置为d_model的4倍左右。
class Seq2SeqTransformer(nn.Module): def __init__(self, vocab_size, d_model=64, nhead=4, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=256, max_len=100): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model) self.pos_encoder = PositionalEncoding(d_model, max_len) self.transformer = nn.Transformer( d_model=d_model, nhead=nhead, num_encoder_layers=num_encoder_layers, num_decoder_layers=num_decoder_layers, dim_feedforward=dim_feedforward, batch_first=True, ) self.fc_out = nn.Linear(d_model, vocab_size) def forward(self, src, tgt, tgt_mask=None, src_pad_mask=None, tgt_pad_mask=None): src = self.pos_encoder(self.embedding(src)) tgt = self.pos_encoder(self.embedding(tgt)) memory = self.transformer.encoder(src, src_key_padding_mask=src_pad_mask) out = self.transformer.decoder(tgt, memory, tgt_mask=tgt_mask, tgt_key_padding_mask=tgt_pad_mask) return self.fc_out(out)这里必须注意batch_first=True。之所以显式设置，是因为nn.Transformer的默认行为让人非常困惑——输入维度是(seq_len, batch, embed_dim)，如果你按照平时熟悉的(batch, seq_len)直觉去构造数据，第一个维度就会完全错位，而且报错信息往往要到矩阵乘法阶段才会出现，明面上看不出问题。
3.3 训练循环与关键参数配置
训练这段代码有几个细节值得说。首先是生成目标序列的mask，这里用nn.Transformer自带的静态方法generate_square_subsequent_mask，它会生成一个右上角为负无穷的方阵。计算loss时要忽略padding位置的损失，用CrossEntropyLoss的ignore_index参数直接指定PAD对应的token id。
学习率上，我试过几组配置，1e-3配合Adam在这个任务上是比较稳的，再往上就容易出现loss震荡甚至直接NaN。一个我在多次实验里验证过的小经验是：Transformer对学习率比较敏感，尤其是输入数据比较稀疏、Embedding表比较大的时候，学习率超过一定阈值，模型数值很容易崩。
model = Seq2SeqTransformer(vocab_size=20) criterion = nn.CrossEntropyLoss(ignore_index=PAD) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for epoch in range(30): model.train() total_loss = 0.0 for step in range(500): src, tgt_input, tgt_output = generate_batch( batch_size=32, vocab_size=20, max_len=8 ) tgt_mask = nn.Transformer.generate_square_subsequent_mask( tgt_input.size(1) ) # pad位置标记为True src_pad_mask = (src == PAD) tgt_pad_mask = (tgt_input == PAD) logits = model(src, tgt_input, tgt_mask=tgt_mask, src_pad_mask=src_pad_mask, tgt_pad_mask=tgt_pad_mask) loss = criterion(logits.reshape(-1, logits.size(-1)), tgt_output.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"Epoch {epoch+1:02d} | Loss {total_loss/500:.4f}")这个训练过程在普通笔记本CPU上大概两三分钟就能收尾。loss会从最初的2.5左右（因为20个类别的随机预测，交叉熵约等于ln20≈3，加上任务学习难度，初始值差不多这个量级）一路降到接近0.01甚至更低。要注意的是，loss降得低不代表模型真的学好了，关键要看推理结果是否和期望完全一致。
3.4 推理与效果验证
推理和训练最大的区别是：训练阶段有teacher forcing，模型每一步都喂真实token；推理阶段没有，你必须把上一步生成的token作为下一步的输入，这就是自回归解码。这个过程如果不写循环，很多人会以为直接把模型forward一次就能出结果，那是不对的。
def greedy_decode(model, src, max_len=20): model.eval() src = src.unsqueeze(0) # (1, seq_len) src_pad_mask = (src == PAD) memory = model.transformer.encoder( model.pos_encoder(model.embedding(src)), src_key_padding_mask=src_pad_mask ) ys = torch.full((1, 1), SOS, dtype=torch.long) for _ in range(max_len): tgt_mask = nn.Transformer.generate_square_subsequent_mask(ys.size(1)) out = model.transformer.decoder( model.pos_encoder(model.embedding(ys)), memory, tgt_mask=tgt_mask ) logits = model.fc_out(out[:, -1, :]) # 只取最后一个位置 next_token = logits.argmax(dim=-1) ys = torch.cat([ys, next_token.unsqueeze(0)], dim=1) if next_token.item() == EOS: break return ys[0] # 去掉batch维度 # 验证一个具体例子 src_example = torch.tensor([3, 5, 1, 4, 2]) result = greedy_decode(model, src_example) print("输入序列:", src_example.tolist()) print("倒序结果:", [t.item() for t in result[1:-1]]) # 去掉SOS和EOS推理代码里有一个细节：位置编码用的是输入对应的最大长度切片，也就是ys的当前长度。每多生成一个token，位置编码就要多取一个位置，因此把位置编码放进循环里每次重新计算是必要的。如果为了省事把位置编码固定成最大长度，效果不会变，但那种写法复用的边界条件太多，容易绊倒人。
我实际跑的结果是，模型在长度为8以内的序列上几乎100%正确。把它放到长度12的序列上也能稳定输出倒序，说明模型确实学到了通用规律，而不是死记硬背。这是我觉得这个小任务成功的标志——它验证了Transformer的序列建模能力。
4. 常见问题与排查技巧实录
这一节是我实际调这个示例过程中踩过的坑，以及一些用户在复现时反馈过的问题整理。我把它们按症状分类，整理成速查表，方便你遇到类似问题时直接定位。
4.1 维度不匹配问题速查
PyTorch里Transformer相关的维度报错，出现的频率远高于正常期望。常见的几种情况分别对应不同原因。如果报错信息出现在MatMul相关的代码行，通常是d_model和nhead的整除关系出了问题——比如d_model=256、nhead=3，256除以3的余数不是0，reshape时报错。如果报错信息出现在Embedding层之后，通常是输入token id的数值超出了词表范围——比如词表大小设为20，但数据里生成了等于20或更大的id，Embedding查表时会直接索引越界。如果报错信息在编码器输出和解码器交互那段，大概率是batch_first设置不一致导致的维度顺序混乱。
4.2 Mask用错导致模型“偷看未来”
这个坑是Transformer训练中最隐蔽的。训练时忘了加tgt_mask，loss照样会降，而且降得还挺快——因为模型可以通过后面的位置直接读取答案。但一到推理阶段，没有teacher forcing提供正确答案，模型立刻露出原形，输出乱七八糟。
判断你有没有踩这个坑有个简单方法：观察训练era的loss。正常训练时，loss下降曲线是平滑的，最终趋于很小的值。如果模型“作弊”，loss会下降得异常快，用两三个epoch就逼近0。然后你在推理时发现准确率远低于训练时的预期。真遇到了也不慌，加上tgt_mask重新训练就行。
另外还有一类易混淆问题是src_pad_mask和tgt_pad_mask用错了串。前者作用于编码器的输入序列，后者作用于解码器的输入序列。如果你只设置了其中一个而漏了另一个，loss最终停在某个较高水平下不去，因为padding位置在不断贡献无关信息。
4.3 训练不收敛与学习率问题
如果你发现loss一直横盘不动，先不要怀疑模型结构。按这个顺序排查：第一，检查数据里是否有NaN——随机生成的任务里不会出现，但如果以后替换成真实文本数据，很可能是文本清洗不干净导致有非法字符；第二，检查学习率，Transformer对学习率极其敏感，尤其是d_model较小、层数较浅时，3e-4到1e-3通常是安全区间；第三，检查梯度，在loss.backward()之后打印一下model.fc_out.weight.grad的数值分布，如果全是NaN，问题一定出在前向过程或loss计算上。
还有一个不太起眼但实际很坑的现象是：loss降到一定程度后不再下降，一直保持在0.5左右震荡。这种情况大概率是位置编码加错了地方——我遇到过有人把位置编码加在了nn.Transformer之后而不是Embedding之后，相当于把位置信息混进注意力输出里，模型很难从中学到有效规律。
4.4 推理阶段的自回归坑点
推理时的坑点比训练更隐蔽。第一个问题是无限循环——如果解码器没有遇到EOS，while循环就会一直跑下去。所以greedy decode函数里一定需要max_len上限保护，到了上限即使没出EOS也要强制截断。第二个问题是生成结果长度比预期少一个——很多人忘了解码器输入要带SOS，导致预测结果整体向右错了一位，结果少了最后一个EOS。
这两个坑在工程上尤其常见。做实时翻译系统时，如果解码循环没有长度上限，线上服务可能因为一个异常样本直接挂掉。即使你的任务没有EOS概念（比如文本分类），也要避免写死循环，必须设置安全上界。
4.5 从小示例到真实任务的三个建议
最后说点题外话。跑通这个倒序示例之后，想往真正的任务迁移，有三个方向可以走。
第一，替换数据源。参考当前开源社区里常见的方式，把一个语料库的句子对整理成(src, tgt)格式，用分词器把句子转成token id，注意设置合理的词表大小和截断长度，然后复用这个训练循环即可。第二，引入更高级的解码策略。greedy decode虽然简单，但实际效果往往不是最好。你可以在next_token = logits.argmax(dim=-1)这里替换成采样、top-k采样或者带温度系数的softmax采样，都能明显改善生成结果的多样性。第三，加一些工程优化。比如混合精度训练、梯度累积和余弦退火学习率，这些都可能在你的场景里带来进一步的收益。
我在实际使用中发现，很多人做完这类小示例后，最大的收获不是学会了调用API，而是终于搞懂了mask到底分几种、维度应该怎么管理、训练和推理的差异在哪里。这些如果只靠听课记笔记，是很难产生真实体感的。倒序任务虽然小，但它像一台精准的解剖台，把Transformer的每个肌理都摊开在你面前。建议你把这段代码保存在本地，之后遇到任何关于Transformer结构的问题，都可以拉出来改一改跑一遍，比翻文档快得多。
