从输入输出切入手撕Transformer：PyTorch代码实现与避坑指南-北京尧图网络

在动手写Transformer的代码之前我踩过最大的一个坑是模型结构背得滚瓜烂熟但真正把数据喂进去的那一刻却被张量维度、位置编码的叠加方式、mask的形状这些边角料卡了整整两天。后来我才想明白Transformer难的不是自注意力的公式本身而是输入怎么变成模型能吃的张量、输出又怎么变回人能看懂的token这一整条链路。这篇就来把Transformer输入输出这条主干道彻底拆开用PyTorch一行行实现给你看。不管你是刚学完注意力机制想动手跑通第一版模型的新手还是已经会调库但没真正手写过底层的新同学这篇都能让你对数据进、结果出这件事有清晰的掌控感Transformer、PyTorch、代码实现三件事一次讲透。1. 为什么值得从输入输出切入手撕Transformer很多人学Transformer的路径是先啃论文里的缩放点积注意力公式推导完softmax和除根号dk然后一头扎进多头注意力的矩阵拆分里最后被残差、层归一化的顺序搞晕。我不是说这条路径不对而是它的门槛设得太高——你还没搞清楚模型吃什么就先研究它怎么嚼挫败感会很强。1.1 输入输出是整条链路的骨架Transformer本质是一个序列到序列的映射函数。输入端要做的事情只有三件把离散的token变成连续向量、给向量注入位置信息、把不同长度的序列补齐对齐以便批量计算。输出端要做的事情也只有两件把每个位置的隐状态投影回词表空间、用交叉熵把预测分布和真实标签对齐。你把这几件事拆清楚中间那堆注意力层其实就是个黑盒哪怕你不完全理解每个矩阵乘法也能先把模型跑起来、把loss降下去。我个人的经验是先跑通输入输出、看到模型真的在学再去深挖注意力的细节学习动力会强很多。因为你能实时看到loss下降、看到生成的文本从乱码逐渐变成通顺的句子这种正向反馈比单纯看公式有用得多。1.2 输入输出藏着最多看起来对其实错的坑真正让我栽跟头的全是输入输出层面的细节。比如位置编码到底是在embedding之后加还是之前加、padding的mask方向和注意力mask的方向是否搞反、训练时输出要不要shift一位做teacher forcing、推理时的自回归循环怎么保持维度一致。这些坑有个共同特点代码能跑通、不报错、loss也在动但结果就是不对。你如果没系统拆过输入输出debug时根本无从下手。所以我建议的路线是把输入输出的每一行都写清楚每个张量的形状都标注出来每步操作的意图都记下来。这样当模型表现异常时你可以按数据进→编码→注意力→输出→损失这条线逐个环节排查而不是对着报错发呆。这份掌控感才是手撕的价值所在。1.3 一套代码覆盖三类典型任务输入输出这条链路设计得好是可以复用的。文本分类任务里你取输出序列的特殊位置向量接个分类头机器翻译任务里你靠输出序列做自回归解码语言模型任务里你把输出和输入错开一位做下一词预测。三者共用同一套embedding、位置编码、mask构建逻辑区别只在最后的输出头和损失函数。理解了这一点你写一次底层就能套用到大半的NLP任务上性价比极高。2. 输入端三大组件逐层拆解输入这条链路看着简单其实每一环都有讲究。我把它拆成embedding、位置编码、mask三块逐块讲清楚为什么这么设计、参数怎么选、代码怎么写。2.1 Token Embedding词表到向量的第一跳embedding层干的事情本质是一张查找表把每个token的id映射成一个d_model维的稠密向量。词表大小vocab_size和模型维度d_model是两个必须提前定下来的超参数。我见过不少新手把vocab_size设得过小导致大量token被迫映射到同一个unknown上信息从一开始就丢了也有人把d_model设成19、37这种奇数后面拆多头的时候直接除不尽。常见的配置是d_model取64、128、256、512、768这些2的幂或者4的倍数这样除以注意力头数nhead能得到整数。比如d_model512、nhead8每个头的维度就是64干净利落。vocab_size一般根据你的分词器实际词表来定中英文混合场景下几万到十几万都正常。我的建议是先用小词表小维度快速验证流程通不通确认没问题再放大否则一上来就上大模型一个epoch跑半小时调试效率极低。还有一个容易被忽视的细节embedding层的初始化。PyTorch默认用正态分布初始化均值0标准差1。但Transformer论文里提到embedding要乘上一个缩放因子也就是根号下d_model。这一步很多教程都一笔带过其实很关键。因为位置编码的数值范围大约在[-1, 1]之间而未经缩放的embedding数值方差接近1两者量级差距大直接相加会让位置信息被淹没。乘以根号d_model后embedding的量级和位置编码匹配两者才能有效融合。import torch import torch.nn as nn import math class TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model, padding_idx0): super().__init__() self.d_model d_model self.embed nn.Embedding(vocab_size, d_model, padding_idxpadding_idx) def forward(self, x): # x: [batch, seq_len] # 乘以 sqrt(d_model) 缩放让量级和位置编码对齐 return self.embed(x) * math.sqrt(self.d_model)注意padding_idx这一项一定要设它会让padding对应的embedding向量在训练中保持为0不参与梯度更新。如果你忘了设padding位置会被当成正常token学习序列越长padding占比越大对模型的污染越严重。2.2 位置编码让模型知道谁先谁后Transformer的自注意力机制本身是位置无关的。换句话说你把输入序列打乱顺序喂进去注意力算出来的结果除了位置对调之外完全一样。但语言是有顺序的猫追老鼠和老鼠追猫意思完全不同。所以必须显式地给每个位置注入顺序信息这就是位置编码的作用。原论文用的是正弦余弦位置编码公式是固定的、不需要学习偶数维度PE(pos, 2i) sin(pos / 10000^(2i/d_model))奇数维度PE(pos, 2i1) cos(pos / 10000^(2i/d_model))我第一次看到这个公式时也懵为什么用正弦余弦、为什么底数是10000。用生活类比解释一下这个设计让每个位置在每个维度上都有一个唯一的指纹而且不同位置之间的相对关系可以通过三角恒等式线性表达出来模型只需要学一个线性变换就能从任意位置的编码推算出相对距离。这就是它能泛化到训练时没见过的更长序列的原因。具体实现时有个坑用exp和log的组合来计算分母比直接连乘幂次数值更稳定。下面这份实现我用了很多次实测很稳。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(pdropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # [max_len, 1] # 用 exp(log) 组合计算避免大指数溢出 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) # [1, max_len, d_model] # register_buffer 让 pe 随模型保存但不参与梯度更新 self.register_buffer(pe, pe) def forward(self, x): # x: [batch, seq_len, d_model] x x self.pe[:, :x.size(1), :] return self.dropout(x)提示pe用register_buffer注册而不是普通属性这样模型保存时会带上它加载时自动恢复同时它不会被优化器当成参数更新。这个细节我一开始没注意后来发现模型换设备时pe没跟着走直接报device不匹配的错。positional encoding还有一个变体是可学习的位置嵌入直接建一个[max_len, d_model]的embedding表让模型自己学。两种方案各有取舍固定编码泛化性好、参数少适合序列长度不固定的场景可学习编码更灵活但需要见过足够长的序列才能学好对应的位置。我个人在序列长度比较固定的任务上更倾向可学习的方式泛化要求高的场景用正弦编码。2.3 归一化与Dropout的接入位置输入端还有两个容易被忽略但影响很大的组件层归一化和Dropout。原始论文里层归一化放在每个子层之后也就是Post-LN结构残差相加后再归一化。但后来大量实践发现Pre-LN更稳定也就是先归一化再进子层。Post-LN的问题在于残差相加后的值会随着层数增加不断累积导致深层梯度不稳定训练时需要精心的warmup策略。Pre-LN则把归一化提前主干的数值范围更可控训练收敛快很多对学习率也没那么敏感。我现在的习惯是一律用Pre-LNPyTorch从1.12开始TransformerEncoderLayer就支持norm_first参数直接设成True即可。Dropout的放置也有讲究。输入端embedding和位置编码相加后会加一次dropout每个注意力子层的输出后加一次前馈网络的中间层也加一次。dropout率一般取0.1小数据量时可以提到0.2到0.3防过拟合。手写的时候记得所有dropout层的训练和推理状态切换要一致用model.train()和model.eval()统一控制别手动去设p值。3. 屏蔽掩码的构建细节mask是Transformer里最容易被写错、又最难察觉的地方。它不控制计算只控制注意力看哪里所以错了也不报错只是结果悄悄变差。3.1 Padding Mask与Causal Mask的区别Transformer里常见的mask有两类用途完全不同千万别混。padding mask解决的是批量内序列长度不一的问题。短句补了padding这些padding位置没有语义其他位置不应该关注它们所以要把注意力分数置成负无穷softmax后自然变成0。它的形状通常是[batch, seq_len]标记每个位置是不是padding。causal mask解决的是自回归生成时的未来信息泄露问题。预测第t个词时只能看到前t个词不能偷看后面的答案。它是个上三角矩阵右上部分置负无穷把未来位置挡住。形状是[seq_len, seq_len]每条下三角为True。我踩过的坑就是这两个mask的形状和方向搞混。padding mask在送进attn_mask或key_padding_mask时要转成合适的形状PyTorch的TransformerEncoderLayer用的是src_key_padding_mask形状是[batch, seq_len]而TransformerDecoderLayer用的是tgt_mask形状是[seq_len, seq_len]。两个接口的命名风格不统一很容易传错。def create_padding_mask(seq, pad_id0): # seq: [batch, seq_len] # 返回 [batch, seq_len]True 表示该位置是 padding需要屏蔽 return (seq pad_id) def create_causal_mask(seq_len, device): # 返回 [seq_len, seq_len]True 表示需要屏蔽上三角 return torch.triu( torch.ones(seq_len, seq_len, devicedevice, dtypetorch.bool), diagonal1 )注意PyTorch的mask约定是True的位置被屏蔽。也就是说你标记为True的地方注意力会把它忽略掉。这和我一开始的直觉相反我总想着True是保留结果写反了模型训练很久loss都不降。3.2 mask在注意力中的生效路径理解了mask的语义还要理解它在哪里生效。注意力计算是Q乘K转置得到分数矩阵再除以根号dk然后加mask最后softmax再乘V。mask是在softmax之前加的把需要屏蔽的位置加上一个极大的负数softmax出来后这些位置的权重就接近0。这里有个数值细节加的是负无穷还是负的极大值。理论上是负无穷但代码里直接用float(-inf)可能在某些情况下导致NaN尤其是整行都被mask掉的时候。稳妥的做法是用一个足够大的负数比如负1e9或者用torch.finfo(dtype).min。此外还要注意如果某一行全是masksoftmax后会出现0/0的NaN一般要给padding位置或者加个安全值处理但在正常的自回归和padding场景下不会出现整行全屏蔽所以问题不大。3.3 维度对齐的检查清单我整理了一份mask维度速查每次写新模型都对着核一遍能省下大量debug时间。mask类型典型形状使用接口作用padding mask[batch, seq_len]src_key_padding_mask屏蔽padding位置causal mask[seq_len, seq_len]tgt_mask / attn_mask屏蔽未来位置合并mask[batch, seq_len, seq_len]attn_mask同时屏蔽两类实际写码时我会在forward里加几行assert把关键张量的shape打出来跑第一个batch时确认一遍之后就可以删掉。这个习惯帮我发现了无数次形状对不上的问题强烈建议你也养成。4. 输出端到损失函数的关键步骤输入端搞定了注意力层跑完接下来是输出端。这里有两件事把隐状态投影回词表以及用损失函数把预测和标签对齐。4.1 线性投影与权重共享每个位置的隐状态是个d_model维向量要变成词表上的概率分布需要过一个线性层把维度从d_model映射到vocab_size然后接softmax得到每个词的概率。这一步本身很简单但有个优化技巧值得说权重共享。所谓权重共享就是让输出线性层的权重直接复用embedding层的权重矩阵。理由是embedding做的是词表到向量的映射输出层做的是向量到词表的映射两者互为逆操作共享权重能减少参数量还能让两边的语义空间保持一致实践中往往能带来性能提升。实现上只要一行代码把两个权重绑定即可。self.fc_out nn.Linear(d_model, vocab_size) self.fc_out.weight self.embedding.embed.weight # 权重绑定提示PyTorch的nn.Linear权重形状是[out_features, in_features]也就是[vocab_size, d_model]而nn.Embedding的权重形状是[num_embeddings, embedding_dim]也就是[vocab_size, d_model]两者一致可以直接绑定。绑定后记得embedding层不要设padding_idx的向量为0后又被反传破坏实践上没问题。4.2 teacher forcing与标签移位训练自回归模型时输入序列和标签序列要错开一位。假设句子是[A, B, C, D]输入是[A, B, C]标签是[B, C, D]。模型看到A预测B看到A B预测C以此类推。这就是teacher forcing用真实的上一个词而不是模型自己预测的词作为下一步输入训练更快更稳。代码实现时如果你用的是nn.Transformer输入和输出的embedding是分开处理的。输入过encoder得到记忆输出右移一位后过decoder再投影到词表和标签算交叉熵。这里移位容易出错我用的是切片方式tgt_input tgt[:, :-1]tgt_label tgt[:, 1:]。注意还要给tgt_input前面加个起始符或者直接用切片让decoder从第一个词开始。# 假设 tgt: [batch, seq_len]包含 bos ... eos tgt_input tgt[:, :-1] # 去掉最后一个作为decoder输入 tgt_label tgt[:, 1:] # 去掉第一个作为预测目标交叉熵损失要注意ignore_index参数把padding位置的损失忽略掉否则模型会花大量精力去预测padding。loss算完后除以有效token数做平均而不是除以总长度这样不同批次的loss才可比。criterion nn.CrossEntropyLoss(ignore_index0) # 0是padding id # logits: [batch, seq_len, vocab_size] loss criterion(logits.reshape(-1, vocab_size), tgt_label.reshape(-1))4.3 完整的前向流程串一遍把上面这些拼起来一个最小的Transformer语言模型前向就是这样的class MiniTransformer(nn.Module): def __init__(self, vocab_size, d_model256, nhead8, num_layers4, dim_ff1024, dropout0.1, max_len512): super().__init__() self.d_model d_model self.embedding TokenEmbedding(vocab_size, d_model, padding_idx0) self.pos_enc PositionalEncoding(d_model, max_len, dropout) encoder_layer nn.TransformerEncoderLayer( d_modeld_model, nheadnhead, dim_feedforwarddim_ff, dropoutdropout, batch_firstTrue, norm_firstTrue ) self.encoder nn.TransformerEncoder(encoder_layer, num_layers) self.fc_out nn.Linear(d_model, vocab_size) self.fc_out.weight self.embedding.embed.weight def forward(self, src, src_key_padding_maskNone, attn_maskNone): # src: [batch, seq_len] x self.embedding(src) # [B, L, D] x self.pos_enc(x) # [B, L, D] x self.encoder(x, maskattn_mask, src_key_padding_masksrc_key_padding_mask) logits self.fc_out(x) # [B, L, vocab] return logits这份代码可以直接跑通一个字符级语言模型。我建议你拿一段中文文本、按字符建词表跑上几十个epoch看着它从生成乱码慢慢变成能拼出词这个过程对理解整条链路帮助极大。5. 训练与推理阶段的输入输出差异训练和推理看起来只是调用方式不同实际上输入输出的组织方式差别很大。搞不清这一点很容易写出训练正常但推理拉胯的模型。5.1 训练阶段的并行与推理阶段的串行训练时因为有teacher forcing整个目标序列可以一次性喂进去并行计算所以Transformer训练效率高。但推理时你没有真实标签只能一个词一个词地生成生成第t个词需要前面t-1个词的结果天然是串行的。这就导致推理速度远慢于训练序列越长越明显。如果你的任务是纯编码类比如文本分类、序列标注那推理也是并行的没这个问题。但生成类任务就要注意了。5.2 自回归推理循环的实现要点自回归推理的核心循环是维护一个已生成序列每次把整个序列喂进去前向一次取最后一个位置的logits选出下一个词追加到序列末尾直到生成结束符或达到最大长度。torch.no_grad() def generate(model, start_tokens, max_new_tokens50, temperature1.0, top_kNone): model.eval() tokens start_tokens # [1, seq_len] for _ in range(max_new_tokens): logits model(tokens) # [1, L, vocab] next_logits logits[:, -1, :] / temperature # 取最后一个位置 if top_k is not None: v, _ torch.topk(next_logits, top_k) next_logits[next_logits v[:, [-1]]] float(-inf) probs torch.softmax(next_logits, dim-1) next_token torch.multinomial(probs, num_samples1) # [1, 1] tokens torch.cat([tokens, next_token], dim1) if next_token.item() eos_id: break return tokens注意这个朴素实现每步都要把整个序列重新前向一遍重复计算极其严重序列长到几百时慢得离谱。优化思路是KV Cache——把已经算过的key和value缓存下来每步只对新token计算这是所有推理加速方案的基础。手写阶段可以先不优化跑通逻辑要紧但要知道这个瓶颈在哪。温度参数控制生成的随机性temperature小于1让分布更尖锐、生成更确定大于1让分布更平坦、生成更多样。top_k则是只在概率最高的前k个词里采样避免采到长尾怪词。这两个是生成任务最常用的调节旋钮建议都留出接口。5.3 一个完整的训练循环模板把前面的东西拼成一个训练循环结构大概是这样def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0.0 for batch in loader: src batch[input_ids].to(device) # [B, L] tgt batch[labels].to(device) pad_mask create_padding_mask(src, pad_id0) logits model(src, src_key_padding_maskpad_mask) loss criterion(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1)) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() total_loss loss.item() return total_loss / len(loader)梯度裁剪这一步我单独拎出来说Transformer对梯度爆炸比较敏感尤其是学习率偏大的时候。clip_grad_norm_把梯度的整体范数限制在max_norm以内是防止loss突然飞掉的保险丝。我在没加这行之前遇到过loss从1.2直接跳到nan的情况加上之后就稳了。6. 常见问题与排查技巧实录下面这些是我和身边朋友手撕Transformer时反复遇到的真问题整理成速查表遇到卡壳时对着看。现象可能原因排查方法loss几乎不下降mask方向写反或学习率过大检查mask的True/False语义打印一个batch的注意力权重看屏蔽对不对loss降了但生成乱码训练和推理的标签移位不一致对比训练时的tgt_input和推理时的起始输入换设备报device不匹配位置编码没注册成buffer用register_buffer确认pe和输入在同一device长序列结果差训练时max_len不够位置编码越界打印最大序列长度确认没超PositionalEncoding的max_len内存爆掉batch和序列长度太大减小batch或改用梯度累积模拟大批量维度报错d_model除不尽nhead确保d_model % nhead 0我重点讲几个不是一眼能看出来的。第一个是mask写反。PyTorch的约定是True表示屏蔽但很多人习惯用True表示保留赋值时反着来。这个错误的隐蔽性在于代码不报错甚至loss也能缓慢下降只是永远到不了理想水平。我的排查手段是构造一个极端的短序列手动把mask打印出来用眼睛核对屏蔽区域是不是我想要的。构造小样本验证是个万能技巧遇到任何可疑行为都先缩小到能一眼看穿的最小例子。第二个是标签位移。训练时我习惯tgt_input tgt[:, :-1]、tgt_label tgt[:, 1:]但推理时起始输入只给了起始符模型看到的信息比训练时少。如果训练时序列开头没放起始符模型会不适应。解决办法是训练和推理都统一带上起始符保证分布一致。这类训练和推理不一致的问题在工程里非常常见本质是数据管道的两套代码走了不同逻辑。第三个是位置编码越界。PositionalEncoding的pe表大小是max_len如果你喂进去的序列比它长切片时虽然不会报错但超出的位置拿不到位置编码等于没位置信息。我建议在forward里加个长度断言超出就报警或动态扩展别让它静默出错。实操心得每次写完一个新模型先拿两三条假数据过一遍forward把每个中间张量的形状打出来核对。这个动作花不了两分钟但能挡掉八成以上的低级错误。我早期省了这一步结果花了几个小时在错误信息里绕圈。最后一个经验是学习率调度。Transformer原论文用的是warmup加逆平方根衰减先把学习率线性升上去再缓慢降下来。这个策略对深层Transformer很关键warmup让模型在初期参数还乱的时候小步走避免一开始就把参数带偏。我一般用固定步数warmupwarmup步数占总步数的5%到10%之后用余弦退火或线性衰减。PyTorch的LambdaLR可以很方便地实现这套调度配合前面说的梯度裁剪训练稳定性会好很多。把输入输出这条链路彻底拆透之后我再去看那些注意力公式、多头拆分、前馈网络就都是可以被顺畅嵌入到这条主干上的零件了。真正决定模型能不能跑起来、跑得对不对的往往不是那些精妙的矩阵变换而是embedding有没有缩放、mask有没有写反、标签有没有移位这些不起眼的地方。我现在的习惯是每换一个新任务先把数据管道的输入输出打印出来逐字段核对确认这条链路干净了再往中间堆结构。这个顺序反过来debug的代价会大得多。

从输入输出切入手撕Transformer：PyTorch代码实现与避坑指南

相关推荐

【小白也能轻松用】最新版OpenClaw v2.7.9配置方法：Windows下用TaoToken统一Key接入数字员工（含最新安装包）

Java反射和注解

软考架构设计师论文 —— 论AI驱动的智能运维架构及其应用（3）

更多资讯

CodeBuddy + WorkBuddy 实战：AI IDE 与 Agent 工作台如何打通开发全链路

Model-Optimizer：面向AI工程落地的模型交付决策框架

大数据入门实战：Linux操作与Hadoop伪分布式搭建实验指南

基于YOLOv11的无人机电力设备异常检测与定位系统设计

gpt-image-1生产实践：蒙版、Alpha通道与透明背景生成

双塔模型：推荐系统中解耦用户与物品表征的工业级召回范式

今日更新

MATLAB+Yalmip+CPLEX实战：综合能源系统优化调度全流程解析

I3C比I2C快10倍？RK3576实战：速率、DTS配置与混合总线避坑指南

字符串转对象：JSON.parse、new Function与URLSearchParams

本周精选

如何彻底离线使用 Laya-CoreML：Hugging Face 下载、缓存机制与 local_files_only 详解

基于1Panel 部署MySQL MCP 服务：TaoToken 统一 Key 接入与 config.toml 配置实战

崇州园区营销网站建设避坑指南：5大注意事项保备案安全

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？