我一直有个习惯凡是热门模型一定要用TensorFlow 2.0亲手实现一遍才踏实。Transformer就是其中之一。说实话用官方tf.keras.layers.MultiHeadAttention几行就能搭完一个encoder但那只是“会用”不是“理解”。这篇文章我从注意力机制的核心拆起把位置编码、多头注意力、掩码、Encoder和Decoder全部手写一遍代码可以直接跑适合已经会Python、想深入掌握Transformer原理并能独立实现的读者。1. 动手前先理清Transformer的四个核心组件网上关于Transformer的解释很多但真正要动手写代码之前很多细节必须自己先想明白否则写出来就是一堆形状不对的报错。我建议先把下面这四个问题搞定再开始敲代码。1.1 自注意力机制到底在算什么自注意力Self-Attention本质上是在计算“序列中每个位置相对于其他所有位置的权重”。对于输入序列的每一个token向量它会生成三个向量Query查询、Key键、Value值。你可以把它们类比成图书馆检索场景Query是你的“检索需求”Key是书的“标签”Value是“书的内容”。通过Query和所有Key算相似度得到注意力权重再用权重把Value加权求和就是当前token融合了整个序列信息后的新表示。在代码层面这个计算就是[ Attention(Q,K,V)softmax(\frac{QK^T}{\sqrt{d_k}})V ]这里除以(\sqrt{d_k})是为了防止点积过大导致softmax梯度消失。(d_k)是key的维度通常等于d_model / num_heads。我在实现时特别注意这一点缩放因子是必须的不是可选项。当年第一次自己写attention忘了除以根号结果训练到第几个epoch损失就NaN了后来才发现是注意力分数太大softmax输出饱和梯度直接崩掉。1.2 多头注意力为什么要“头”多头注意力就是把d_model维度的Q、K、V分成num_heads份每一份独立做自注意力最后再拼回来。这样做的意义不是简单地增加参数而是让每个“头”有机会关注不同类型的关系。比如一个头可能更关注语义相关性另一个头可能更关注位置邻近性。实现时最常见的坑是维度变换顺序。假设batch_sizeB、seq_lenS、d_modelD、head_numH、每个头的维度depthD//H。先从(B,S,D)线性投影到(B,S,D)然后reshape成(B,S,H,depth)但注意下一步必须transpose成(B,H,S,depth)把“头”维度提前。很多人直接reshape后去算attention结果把不同头的向量混在一起模型怎么训都学不好。1.3 位置编码给序列注入顺序信息RNN靠逐步输入自带顺序CNN靠卷积感受野也能感知局部顺序而Transformer是纯注意力本身对“位置”一无所知。你把输入序列倒过来注意力计算的结果一模一样。所以必须显式加入位置信息。原论文用的是三角位置编码[ PE(pos, 2i) sin(pos / 10000^{2i/d_{model}}) ] [ PE(pos, 2i1) cos(pos / 10000^{2i/d_{model}}) ]其中(pos)是token下标(i)是维度下标。这里用sin/cos交错有一个巧妙之处任何位置的向量都可以表示成其他两个位置向量、的线性组合这样模型更容易学到相对位置关系。实现的时候我一般直接预计算一个[max_len, d_model]矩阵加到Embedding输出上。注意这里的加不是拼接是逐元素相加所以Embedding的输出维度和位置编码维度必须完全一致。1.4 掩码训练时的“上帝视角”必须被遮住Transformer有两种掩码一种是padding mask用来忽略pad位置另一种是look-ahead mask也叫因果掩码让Decoder在预测第(t)个token时只能看到前(t-1)个token。这个很好理解如果你预测下一个词时能看到答案那还学什么Look-ahead mask是一个上三角矩阵上三角区域置为1屏蔽下三角区域置为0保留。在代码里我会把它做成一个[1, 1, seq_len, seq_len]的布尔矩阵方便和注意力分数广播。2. TensorFlow 2.0环境准备与数据构造模型原理算是捋清楚了接下来是实操。我先说明一下我的运行环境TensorFlow 2.4以上都可以Python 3.7最好有GPU没有GPU用CPU跑一个极小的数据集也能验证代码正确性。2.1 环境版本与依赖pip install tensorflow2.10.0 tensorflow-text不需要安装额外的NLP库分词我们直接用最简单的字符级/单词级切分。如果你的tf版本低于2.0建议直接换新环境因为很多Keras API在1.x里行为完全不同。2.2 构造一个极小的翻译任务数据集为了把原理跑通我造了一个“英文单词→中文单词/短语”的迷你数据集。不用真实语料不然预处理就够写半篇文章。# 模拟数据对 pairs [ (hello, 你好), (how are you, 你好吗), (i love you, 我爱你), (what is this, 这是什么), (good morning, 早上好), (good night, 晚安), ]真正的模型训练肯定需要大量数据但这里我们关心的是实现结构和代码逻辑小数据集够用了。需要说明的是这个数据集非常小最终效果不会好但能证明整个前向/反向/推理流程是通顺的。2.3 数据预处理与tf.data管道我做了两个词表一个英文词表一个中文词表。中文按字符切分英文按空格切分。每个句子前后加start和end并做padding到固定长度。import tensorflow as tf # 构建词表 def build_vocab(sentences, tokenize_fn): vocab {pad: 0, start: 1, end: 2, unk: 3} for s in sentences: for token in tokenize_fn(s): if token not in vocab: vocab[token] len(vocab) return vocab en_vocab build_vocab([p[0] for p in pairs], lambda s: s.split()) zh_vocab build_vocab([p[1] for p in pairs], list)然后把文本转成ID序列并padding到固定长度。tf.data.Dataset的好处是能自动处理batch、shuffle和prefetch。这里有个小经验prefetch(tf.data.AUTOTUNE)一定要加否则数据加载会拖慢GPU训练。def encode(en, zh, en_vocab, zh_vocab, max_len_en10, max_len_zh10): en_tokens [start] en.split() [end] zh_tokens [start] list(zh) [end] en_ids [en_vocab.get(t, en_vocab[unk]) for t in en_tokens] zh_ids [zh_vocab.get(t, zh_vocab[unk]) for t in zh_tokens] en_ids en_ids[:max_len_en] [0] * (max_len_en - len(en_ids[:max_len_en])) zh_ids zh_ids[:max_len_zh] [0] * (max_len_zh - len(zh_ids[:max_len_zh])) return en_ids, zh_idsDecoder的输入和输出要做“偏移”训练时Decoder的输入是start 目标序列输出是目标序列 end。这样每个位置都在“用前一个词预测下一个词”。3. 从零实现Transformer模块这是全文最硬核的部分。我尽量把每个模块都贴出完整代码同时解释为什么要这样写。3.1 位置编码层的实现位置编码层不参与训练直接是一个静态张量。但为了方便调试我把它封装成Keras层输出维度是(batch, seq_len, d_model)。class PositionalEncoding(tf.keras.layers.Layer): def __init__(self, d_model, max_len1000): super().__init__() self.d_model d_model self.max_len max_len pe np.zeros((max_len, d_model), dtypenp.float32) position np.arange(max_len)[:, None] div_term np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model)) pe[:, 0::2] np.sin(position * div_term) pe[:, 1::2] np.cos(position * div_term) self.pe tf.constant(pe[None, ...]) def call(self, x): seq_len tf.shape(x)[1] return x self.pe[:, :seq_len, :]注意div_term的计算原论文的写法是1 / 10000^(2i/d_model)取对数后就是(-2i/d_model) * log(10000)而2i是偶数索引我直接用了np.arange(0, d_model, 2)所以指数是exp(arange * (-log(10000)/d_model))和公式等价。3.2 多头注意力层从Linear到组合先把缩放点积注意力做出来。这个函数接受Q、K、V和mask返回加权后的结果和注意力权重。def scaled_dot_product_attention(q, k, v, maskNone): d_k tf.cast(tf.shape(k)[-1], tf.float32) scores tf.matmul(q, k, transpose_bTrue) / tf.sqrt(d_k) if mask is not None: scores mask * -1e9 # mask中True的位置被屏蔽 weights tf.softmax(scores, axis-1) output tf.matmul(weights, v) return output, weights这里有个容易踩坑的地方掩码和-1e9相加而不是相乘。很多初学者的写法是scores tf.where(mask, -1e9, scores)也可以但mask必须是布尔张量。我更习惯把mask直接变成0保留或1屏蔽的浮点张量这样加一个极大负数就能让softmax概率趋于0。多头注意力层的完整实现如下。class MultiHeadAttention(tf.keras.layers.Layer): def __init__(self, d_model, num_heads): super().__init__() self.d_model d_model self.num_heads num_heads self.depth d_model // num_heads self.wq tf.keras.layers.Dense(d_model) self.wk tf.keras.layers.Dense(d_model) self.wv tf.keras.layers.Dense(d_model) self.dense tf.keras.layers.Dense(d_model) def split_heads(self, x, batch_size): x tf.reshape(x, (batch_size, -1, self.num_heads, self.depth)) return tf.transpose(x, perm[0, 2, 1, 3]) def call(self, q, k, v, maskNone): batch_size tf.shape(q)[0] q self.split_heads(self.wq(q), batch_size) k self.split_heads(self.wk(k), batch_size) v self.split_heads(self.wv(v), batch_size) attn_output, _ scaled_dot_product_attention(q, k, v, mask) attn_output tf.transpose(attn_output, perm[0, 2, 1, 3]) concat_attention tf.reshape(attn_output, (batch_size, -1, self.d_model)) output self.dense(concat_attention) return output为什么Dense(d_model)之后又要reshape又要transpose因为原论文要求每个头独立计算如果不把heads放到batch维度旁边而是保留在seq_len后面attention分数矩阵会变成(B, S, H, S, depth)不仅多一维计算顺序也可能乱掉。transpose之后变成(B, H, S, depth)matmul时Q和K都在最后两维做矩阵乘法就直接完成了所有头的并行计算。3.3 前馈网络与残差层归一化每个Encoder/Decoder子层后面都跟一个前馈网络Feed-Forward Network一般叫FFN[ FFN(x) \max(0, xW_1b_1)W_2 b_2 ]第一层把维度放大4倍第二层再映射回d_model。这个结构很像传统的全连接网络但它的作用是对每个位置独立做非线性变换不跨位置。class FeedForward(tf.keras.layers.Layer): def __init__(self, d_model, dff): super().__init__() self.dense1 tf.keras.layers.Dense(dff, activationrelu) self.dense2 tf.keras.layers.Dense(d_model) def call(self, x): return self.dense2(self.dense1(x))子层连接使用的是Add Norm。Transformer论文使用的是先Add后LayerNorm也就是LayerNorm(x Sublayer(x))而不是Post-Norm。这个顺序会影响训练的稳定性。我按原论文实现class EncoderLayer(tf.keras.layers.Layer): def __init__(self, d_model, num_heads, dff): super().__init__() self.attn MultiHeadAttention(d_model, num_heads) self.ffn FeedForward(d_model, dff) self.layernorm1 tf.keras.layers.LayerNormalization(epsilon1e-6) self.layernorm2 tf.keras.layers.LayerNormalization(epsilon1e-6) def call(self, x, trainingFalse, maskNone): attn_output self.attn(x, x, x, mask) out1 self.layernorm1(x attn_output) ffn_output self.ffn(out1) out2 self.layernorm2(out1 ffn_output) return out2这里training参数好像没用到但留着有好处以后要加Dropout或正则化可以直接用。3.4 Encoder层与Decoder层的拼装Encoder层就是多个EncoderLayer叠加外面再套一个输入Embedding和位置编码。class Encoder(tf.keras.layers.Layer): def __init__(self, vocab_size, d_model, num_heads, dff, num_layers): super().__init__() self.d_model d_model self.embedding tf.keras.layers.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model) self.enc_layers [EncoderLayer(d_model, num_heads, dff) for _ in range(num_layers)] def call(self, x, maskNone): seq_len tf.shape(x)[1] x self.embedding(x) x * tf.sqrt(tf.cast(self.d_model, tf.float32)) # 原文中的缩放 x self.pos_encoding(x) for layer in self.enc_layers: x layer(x, maskmask) return xDecoder和Encoder很像但有三个关键区别第一层是masked自注意力使用look-ahead mask第二层是encoder-decoder注意力Query来自Decoder而Key和Value来自Encoder的输出使用的mask是padding mask第三层才是FFN。另外每一层都要保留残差和LayerNorm。class DecoderLayer(tf.keras.layers.Layer): def __init__(self, d_model, num_heads, dff): super().__init__() self.attn1 MultiHeadAttention(d_model, num_heads) self.attn2 MultiHeadAttention(d_model, num_heads) self.ffn FeedForward(d_model, dff) self.layernorm1 tf.keras.layers.LayerNormalization(epsilon1e-6) self.layernorm2 tf.keras.layers.LayerNormalization(epsilon1e-6) self.layernorm3 tf.keras.layers.LayerNormalization(epsilon1e-6) def call(self, x, enc_output, look_ahead_mask, padding_maskNone): attn1_output self.attn1(x, x, x, look_ahead_mask) out1 self.layernorm1(x attn1_output) attn2_output self.attn2(out1, enc_output, enc_output, padding_mask) out2 self.layernorm2(out1 attn2_output) ffn_output self.ffn(out2) out3 self.layernorm3(out2 ffn_output) return out33.5 完整Transformer模型组装把Encoder和Decoder串起来最后加一个Dense(vocab_size)输出每个token的概率分布。训练时用SparseCategoricalCrossentropy作为损失函数。class Transformer(tf.keras.Model): def __init__(self, en_vocab_size, zh_vocab_size, d_model128, num_heads8, dff512, num_layers4, max_len20): super().__init__() self.encoder Encoder(en_vocab_size, d_model, num_heads, dff, num_layers) self.decoder Decoder(zh_vocab_size, d_model, num_heads, dff, num_layers) self.final_layer tf.keras.layers.Dense(zh_vocab_size) def call(self, inputs, trainingFalse): src, tgt inputs src_mask create_padding_mask(src) tgt_mask create_look_ahead_mask(tgt) tgt_padding_mask create_padding_mask(tgt) combined_mask tf.maximum(tgt_mask, tgt_padding_mask) enc_output self.encoder(src, src_mask) dec_output self.decoder(tgt, enc_output, combined_mask, src_mask) return self.final_layer(dec_output)这里有个细节Decoder第一层用的mask是look-ahead mask和padding mask的并集。因为padding位置的token本来就不该被关注同时因果限制也不能看到未来token。我一开始只传了look-ahead结果模型在Decoder输入含padding时会出现奇怪的注意力loss迟迟降不下去。掩码函数实现如下def create_padding_mask(seq): seq tf.cast(tf.math.equal(seq, 0), tf.float32) return seq[:, tf.newaxis, tf.newaxis, :] # (batch, 1, 1, seq_len) def create_look_ahead_mask(seq_len): mask 1 - tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0) return mask[tf.newaxis, tf.newaxis, :, :] # (1, 1, seq_len, seq_len)4. 训练陷阱与调试经验代码写完之后真正折磨人的是训练阶段。这个部分我踩过的坑可以列成一张长长的清单。4.1 学习率预热策略的必要性Transformer通常使用自定义学习率先线性预热再按步数倒数衰减。公式是[ lr d_{model}^{-0.5} \cdot \min(step^{-0.5}, step \cdot warmup_steps^{-1.5}) ]为什么要这么麻烦因为训练初期模型参数是随机的如果直接用大学习率很容易让注意力权重分布过激导致训练崩溃。先用小学习率“预热”几个step让模型稳定下来之后再逐渐衰减。我在Keras中这样做class WarmupSchedule(tf.keras.optimizers.schedules.LearningRateSchedule): def __init__(self, d_model, warmup_steps4000): super().__init__() self.d_model tf.cast(d_model, tf.float32) self.warmup_steps warmup_steps def __call__(self, step): step tf.cast(step, tf.float32) arg1 tf.math.rsqrt(step) arg2 step * (self.warmup_steps ** -1.5) return tf.math.rsqrt(self.d_model) * tf.minimum(arg1, arg2) optimizer tf.keras.optimizers.Adam(WarmupSchedule(d_model128), beta_10.9, beta_20.98, epsilon1e-9)注意Adam的beta_2原论文用的是0.98不是TensorFlow默认的0.999。这个细节对训练稳定性有点影响很多从RNN转过来的资料不会告诉你。4.2 训练时的显存/内存优化Transformer的计算量不小即使小数据集也可能把显存占满。我有几个经验使用tf.data的.batch(...).prefetch(tf.data.AUTOTUNE)不要让GPU等着CPU产生数据。在call方法里使用tf.function加速。在model.compile时Keras会默认开启但自定义train_step时需要手动加tf.function。如果显存不够减小d_model或num_layers不要一开始就上大模型。我调试时甚至用d_model16跑通了再扩大。tf.function def train_step(src, tgt): tgt_input tgt[:, :-1] tgt_output tgt[:, 1:] with tf.GradientTape() as tape: predictions model((src, tgt_input), trainingTrue) loss loss_function(tgt_output, predictions) gradients tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return loss4.3 常见错误维度不匹配、掩码广播问题新手最常遇到的报错就是Incompatible shapes。我整理过一张自查表症状大概率原因解决方案attention scores变成5维split_heads里transpose顺序错确认是(B,H,S,depth)输出shape和输入对不上最终Dense层维度写错Decoder输出维度应为目标词表大小训练loss不降没有做学习率预热换成WarmupSchedule预测时全是end或pad掩码用错模型偷看未来检查Decoder第一层mask梯度爆炸忘了对Embedding做缩放embedding * sqrt(d_model)掩码广播是我花时间最多的地方。注意create_padding_mask返回的shape是(batch, 1, 1, seq_len)而attention scores的shape是(batch, heads, seq_len, seq_len)。(batch,1,1,seq_len)可以广播到(batch,1,seq_len,seq_len)吗不能前者最后一个维度是seq_len用来屏蔽“key序列”中的pad位置所以它的位置应该放在第3维也就是[:, 1, :, 1]。实际上(batch, 1, 1, seq_len)的维度含义是对于每个batch和每个headkey序列位置上的mask值刚好能广播到scores的最后两维中的第二维也就是(batch, heads, seq_len_q, seq_len_k)的seq_len_k维度。这也是我为什么坚持用seq[:, tf.newaxis, tf.newaxis, :]而不是seq[:, tf.newaxis, :, tf.newaxis]的原因。4.4 用teacher forcing还是自回归训练训练Decoder时我们通常使用teacher forcing直接把真实目标序列作为Decoder的每一步输入而不是用模型自己生成的token。这样收敛快但会造成训练/推理不一致exposure bias。解决办法是计划采样scheduled sampling但小项目直接用teacher forcing就行。在推理阶段才改为自回归。TensorFlow里实现teacher forcing很简单上面train_step里tgt_input tgt[:, :-1]就是。但要注意tgt_input和tgt_output长度相同tgt_input是去掉最后一个结束符tgt_output是去掉第一个开始符。5. 推理阶段与效果验证模型训练好之后不能直接调用predict因为predict是喂整个目标序列而推理时我们不知道目标序列。必须一步一步地生成。5.1 贪心解码还是集束搜索最简单的解码策略是贪心每一步选概率最大的token作为当前输出然后把它拼到Decoder输入中继续预测下一个token。代码如下def greedy_decode(model, src, max_len10): start_token zh_vocab[start] end_token zh_vocab[end] encoder_input tf.constant([src]) enc_output model.encoder(encoder_input, create_padding_mask(encoder_input)) decoder_input tf.constant([[start_token]], dtypetf.int64) for _ in range(max_len): combined_mask create_look_ahead_mask(tf.shape(decoder_input)[1]) dec_output model.decoder(decoder_input, enc_output, combined_mask, None) logits model.final_layer(dec_output[:, -1, :]) next_token tf.argmax(logits, axis-1) if next_token.numpy()[0] end_token: break decoder_input tf.concat([decoder_input, next_token], axis-1) return decoder_input这里有个性能问题每次循环Decoder的输入长度都会增加如果max_len很长计算量会越来越大。但小数据集无所谓图个清楚。集束搜索Beam Search能比贪心找到更优的序列但会多消耗几倍显存。对于入门理解建议先跑通贪心再扩展beam search。5.2 让模型跑一次翻译示例训练几十个epoch后模型仍然会因为数据集太小而翻译不准但至少能看到结构是通的。比如输入hello可能输出你好输入i love you也许能生成我爱你。如果输出完全混乱先检查数据切分和词表是否一致。这里多说一句小数据集本身不具备泛化能力所以不要因为效果差就怀疑实现错了。我的验证方法是构造sin/cos位置编码的可视化看每个位置的向量在序列长度上的响应是否平滑再用一个只有几个样本的过拟合测试让模型在训练集上把loss降到接近0如果降不下去说明模型结构有bug。5.3 评估指标BLEU简算机器翻译常用BLEU它的核心是n-gram精确率加上长度惩罚。手工实现非常简单def simple_bleu(reference, candidate): ref_tokens list(reference) cand_tokens list(candidate) matches sum(min(ref_tokens.count(t), cand_tokens.count(t)) for t in set(cand_tokens)) precision matches / len(cand_tokens) if cand_tokens else 0 brevity_penalty min(1.0, len(ref_tokens) / len(cand_tokens)) if cand_tokens else 0 return precision * brevity_penalty真正的BLEU会考虑1-4gram的权重但简版足够你判断“预测的和参考是否接近”。我在调参时就用这个简版作为参考指标而不是盯着loss看因为loss下降不一定代表翻译质量变好。6. 后续扩展与个人感受Transformer实现完成后你会发现很多现代模型都能从这个基础版本演化出来。BERT就是只保留EncoderGPT只保留Decoder而T5、BART则类似Encoder-Decoder结构。理解了这一套再去看那些动辄几十亿参数的预训练模型至少不会觉得陌生。6.1 从Transformer到BERT/GPT的跳跃BERT预训练任务之一是Masked Language Model随机遮住部分token让Encoder预测被遮住的内容。因为Transformer的Encoder是双向注意力的所以天然适合做这种任务。GPT则是Decoder-only结构用因果掩码做自回归语言建模。如果你自己实现的Decoder的masked attention没问题那GPT的基本结构你就已经会了。从代码角度区别只在于mask的设置和Loss计算方式。所以我强烈建议入门Transformer时不要直接使用大模型工具库手写一遍会让后续迁移学习事半功倍。6.2 个人对TensorFlow实现的体会我用TensorFlow 2.0实现Transformer最大的感受是Keras的自定义层和目标API让代码组织很清晰但调试时动态图和tf.function的叠加会有一些诡异问题。一个典型例子是tf.argmax在tf.function里输出维度可能会丢失我在推理时调了半天后来干脆先numpy一下再转tf虽然绕了一点但稳妥。另外model.summary()对这种自定义多输入模型并不友好它经常显示Multi形状未知。我的调试方式是在call方法里临时加print(tf.shape(x))用model((src, tgt))跑一次确认各个模块输出shape正确后再删掉打印。写到这想起自己踩过的最大一个坑训练时loss已经降到接近0但预测时输出全是end。后来发现是训练时Decoder的padding mask和look-ahead mask合并时用了tf.maximum而tf.maximum(0,1)会得到1这没错但tf.maximum(padding_mask, look_ahead_mask)的广播维度不一致导致所有位置都被捂住了。最后我在合并前手动把所有mask都统一成(batch, 1, seq_len, seq_len)问题才消失。如果你也想从头实现Transformer我的建议是先跑通一个极小的玩具翻译任务再尝试加真实数据。整个过程里最值得花时间的就是多头注意力和掩码广播想明白这两个其他都只是堆层数。祝你好运。