1. 从理论到实践为什么需要逐行解读Transformer代码如果你已经看过不少关于Transformer的论文解读和架构图但打开PyTorch实现代码时依然感觉一头雾水那么这篇文章就是为你准备的。理论上的多头自注意力、前馈网络、层归一化在代码里是如何组织成一个高效、可训练的模块的位置编码怎么和词嵌入相加训练和推理时的掩码机制有何不同这些细节光看公式和框图是学不到的。我自己在初次实现Transformer时就踩过不少坑比如忘了在解码器的自注意力层应用前瞻掩码导致模型在训练时“作弊”又比如对层归一化的affine参数理解不透影响了模型收敛。通过这篇超详细的代码解读我希望带你穿透理论直抵工程实现的精髓。我们将基于PyTorch从最基础的张量操作开始一步步搭建一个完整的Transformer模型并解释每一行代码背后的设计意图和数学原理。无论你是想深入理解大模型的基础架构还是准备动手复现或修改Transformer用于自己的任务这篇文章都将提供一份可靠的“地图”。2. 环境搭建与项目结构不仅仅是安装PyTorch在深入代码之前一个清晰、可复现的环境是基石。很多人以为环境搭建就是pip install torch实则不然。2.1 PyTorch版本选择与安装避坑首先访问PyTorch官网获取安装命令是最稳妥的方式。你需要根据你的CUDA版本如果使用GPU来选择。例如对于CUDA 12.1你可能需要安装类似torch2.1.0cu121的版本。一个常见的错误是使用conda安装时默认安装的是CPU版本导致后续无法利用GPU加速。注意如果你遇到类似“invalid archive error”的错误这通常是因为网络问题导致安装包下载不完整。解决方法包括使用国内镜像源如清华源、阿里云源或者直接下载whl文件进行离线安装。我建议创建一个独立的Conda环境来管理项目依赖conda create -n transformer_tutorial python3.9 conda activate transformer_tutorial # 根据你的CUDA版本从官网复制安装命令例如 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121除了PyTorch我们还需要一些辅助库pip install numpy matplotlib tqdmmatplotlib用于可视化注意力权重等tqdm可以给训练循环添加进度条提升体验。2.2 代码文件组织一个清晰的项目结构有助于管理复杂度。建议创建如下目录transformer_code/ ├── model.py # Transformer模型定义核心 ├── train.py # 训练脚本 ├── inference.py # 推理/验证脚本 ├── config.py # 超参数配置模型尺寸、层数等 ├── data_loader.py # 数据加载与预处理 ├── utils.py # 工具函数如位置编码可视化 └── main.py # 主程序入口在model.py中我们将按模块构建Transformer嵌入层、位置编码、多头注意力、前馈网络、编码器层、解码器层最后组装成完整的Encoder-Decoder。这种模块化的方式不仅代码清晰也便于单独测试每个组件。3. 核心模块拆解从嵌入层到注意力机制Transformer的威力源于其精巧的模块设计。我们将自底向上看看每个模块在PyTorch中如何实现。3.1 词嵌入与位置编码给模型注入“顺序”信息Transformer本身不具备处理序列顺序的能力因此需要位置编码Positional Encoding。在代码中这通常是两个独立的层最终相加。词嵌入层Embedding非常简单就是PyTorch的nn.Embeddingimport torch.nn as nn class Embeddings(nn.Module): def __init__(self, vocab_size, d_model): super(Embeddings, self).__init__() self.lut nn.Embedding(vocab_size, d_model) self.d_model d_model # 模型维度如512 def forward(self, x): # x: [batch_size, seq_len] # 输出: [batch_size, seq_len, d_model] return self.lut(x) * math.sqrt(self.d_model) # 缩放参见论文这里乘以math.sqrt(self.d_model)是原论文中的一个小技巧目的是在与位置编码相加前保持嵌入向量的方差在一个合适的范围。位置编码层Positional Encoding的实现则更有趣。它使用正弦和余弦函数来生成绝对位置信息import torch import math class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout0.1, max_len5000): super(PositionalEncoding, self).__init__() self.dropout nn.Dropout(pdropout) # 计算位置编码矩阵形状为 [max_len, d_model] pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # [max_len, 1] div_term torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # 对偶数索引应用sin奇数索引应用cos 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] 便于广播 # 将pe注册为buffer不参与训练但会随模型保存/加载 self.register_buffer(pe, pe) def forward(self, x): # x: [batch_size, seq_len, d_model] x x self.pe[:, :x.size(1)] # 只取前seq_len个位置 return self.dropout(x)关键点在于div_term的计算。公式exp(2i * (-log(10000)/d_model))等价于1/(10000^(2i/d_model))这是原论文中使用的频率项。这样每个位置、每个维度都有一个独特的编码。register_buffer确保这个计算好的pe矩阵能被保存和加载但它不是可训练参数。3.2 缩放点积注意力注意力机制的心脏这是整个Transformer最核心的运算单元。其输入是查询Q、键K、值V三个矩阵。def attention(query, key, value, maskNone, dropoutNone): 计算缩放点积注意力。 参数: query, key, value: 形状均为 [batch_size, num_heads, seq_len, d_k] mask: 可选的掩码形状为 [batch_size, 1, seq_len, seq_len] 或 [batch_size, seq_len, seq_len] dropout: nn.Dropout层实例 返回: 加权后的输出注意力权重 d_k query.size(-1) # 获取键向量的维度 # 计算QK^T并缩放 scores torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: # 将mask中为True或1的位置替换为一个极小的负数这样在softmax后权重接近0 scores scores.masked_fill(mask 0, -1e9) p_attn torch.softmax(scores, dim-1) # 在最后一个维度键序列上做softmax if dropout is not None: p_attn dropout(p_attn) # 用注意力权重加权值向量 return torch.matmul(p_attn, value), p_attn为什么需要缩放除以sqrt(d_k)当d_k较大时点积的结果可能会变得非常大这将softmax函数推入梯度极小的区域导致训练困难梯度消失。缩放操作可以稳定梯度。掩码Mask的作用至关重要。在Transformer中主要有两种掩码填充掩码Padding Mask用于忽略序列中无效的填充位置pad。通常在注意力分数计算前将填充位置对应的分数置为一个很大的负数如-1e9。前瞻掩码Look-ahead Mask仅用于解码器的自注意力层。目的是防止解码器在预测第t个词时“看到”第t个词之后的信息即未来信息。这是一个上三角矩阵对角线及以下为1以上为0。3.3 多头注意力并行化的艺术单头注意力可能只关注到序列中特定模式的依赖关系。多头注意力允许模型同时关注来自不同表示子空间的信息。class MultiHeadedAttention(nn.Module): def __init__(self, h, d_model, dropout0.1): h: 头的数量 d_model: 模型总维度 super(MultiHeadedAttention, self).__init__() assert d_model % h 0 # 确保d_model能被h整除 self.d_k d_model // h # 每个头的维度 self.h h # 定义四个线性层用于生成Q, K, V和最终的输出投影 self.linears nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(4)]) self.attn None # 用于保存注意力权重便于可视化或分析 self.dropout nn.Dropout(pdropout) def forward(self, query, key, value, maskNone): if mask is not None: # 同样的掩码需要应用于所有头 mask mask.unsqueeze(1) # [batch_size, 1, 1, seq_len] 或 [batch_size, 1, seq_len, seq_len] batch_size query.size(0) # 1) 线性投影并分头 # 对Q, K, V分别做线性变换然后重塑形状以分离出头 query, key, value [ lin(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2) for lin, x in zip(self.linears, (query, key, value)) ] # 此时 query, key, value 形状: [batch_size, num_heads, seq_len, d_k] # 2) 在分头后的数据上应用注意力函数 x, self.attn attention(query, key, value, maskmask, dropoutself.dropout) # x形状: [batch_size, num_heads, seq_len, d_k] # 3) 合并多头 x x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) # x形状: [batch_size, seq_len, d_model] # 4) 通过最后一个线性层进行输出投影 return self.linears[-1](x)分头与合并的细节代码中view和transpose的操作顺序是关键。先通过线性层将d_model维的输入映射到d_model维然后通过view重塑为[batch_size, seq_len, num_heads, d_k]再通过transpose(1,2)交换seq_len和num_heads维度得到[batch_size, num_heads, seq_len, d_k]这样每个头就能独立计算注意力。合并时先transpose回来再用view合并。4. 前馈网络与残差连接稳定训练的保障注意力层捕捉的是序列元素间的关系而前馈网络Feed-Forward Network, FFN则对每个位置的表示进行独立的、复杂的非线性变换。4.1 位置级前馈网络FFN是一个简单的两层全连接网络中间有一个ReLU激活函数。class PositionwiseFeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 nn.Linear(d_model, d_ff) # 膨胀层通常d_ff4*d_model self.w_2 nn.Linear(d_ff, d_model) # 收缩层 self.dropout nn.Dropout(dropout) def forward(self, x): # x: [batch_size, seq_len, d_model] # 先膨胀再激活再Dropout最后收缩 return self.w_2(self.dropout(torch.relu(self.w_1(x))))为什么需要这个FFN自注意力层本质上是线性变换的加权和尽管权重是动态计算的。加入带有非线性激活函数如ReLU的FFN极大地增强了模型的表示能力使其能够学习更复杂的模式。4.2 层归一化与残差连接训练深度网络的利器这是Transformer能够成功堆叠多层如12层编码器的关键。每个子层自注意力、FFN都被包裹在一个“残差连接层归一化”的模块中。class SublayerConnection(nn.Module): 一个残差连接后接层归一化。 注意为了简化原论文是先做层归一化再进入子层。但后续研究如Pre-LN Transformer发现先归一化效果可能更好。 这里实现原论文的 Post-LN 结构。 def __init__(self, size, dropout): super(SublayerConnection, self).__init__() self.norm nn.LayerNorm(size) self.dropout nn.Dropout(dropout) def forward(self, x, sublayer): x: 输入张量 sublayer: 一个函数它本身定义了一个子层如多头注意力或FFN # 原论文结构x - LayerNorm - Sublayer - Dropout - Add # 即残差连接应用于子层输出和原始输入x之间 return x self.dropout(sublayer(self.norm(x)))层归一化LayerNorm vs 批归一化BatchNorm对于序列数据BatchNorm会沿着批次维度归一化这在不同序列长度和批次内样本差异大时不稳定。LayerNorm则对每个样本、每个时间步的特征维度进行归一化更适合序列模型。nn.LayerNorm默认启用affine参数即可学习的缩放和偏移这很重要因为它让模型能够学习到最适合当前层的分布。残差连接Residual Connection允许梯度直接流过网络极大地缓解了深度网络中的梯度消失问题。公式output x sublayer(norm(x))是标准做法。5. 编码器与解码器组装核心模块有了上面的积木我们现在可以搭建编码器层和解码器层然后将它们堆叠起来。5.1 编码器层自注意力与FFN的堆叠一个编码器层包含一个多头自注意力子层和一个前馈网络子层每个子层都有残差连接和层归一化。class EncoderLayer(nn.Module): def __init__(self, size, self_attn, feed_forward, dropout): super(EncoderLayer, self).__init__() self.self_attn self_attn self.feed_forward feed_forward self.sublayer nn.ModuleList([SublayerConnection(size, dropout) for _ in range(2)]) self.size size # d_model def forward(self, x, mask): x: 编码器输入 [batch_size, src_seq_len, d_model] mask: 源序列的掩码通常是填充掩码[batch_size, 1, src_seq_len] # 第一个子层多头自注意力自注意力意味着 Q, K, V 都来自同一个输入x x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 第二个子层前馈网络 x self.sublayer[1](x, self.feed_forward) return x注意编码器的自注意力中查询Q、键K、值V都来自上一层编码器的输出对于第一层则是输入嵌入与位置编码的和。mask参数用于屏蔽填充符。5.2 解码器层带掩码的自注意力与编码-解码注意力解码器层比编码器层多一个注意力子层即编码-解码注意力层又称交叉注意力层。class DecoderLayer(nn.Module): def __init__(self, size, self_attn, src_attn, feed_forward, dropout): super(DecoderLayer, self).__init__() self.size size self.self_attn self_attn # 带掩码的自注意力 self.src_attn src_attn # 编码-解码注意力 self.feed_forward feed_forward # 三个子层连接 self.sublayer nn.ModuleList([SublayerConnection(size, dropout) for _ in range(3)]) def forward(self, x, memory, src_mask, tgt_mask): x: 解码器输入 [batch_size, tgt_seq_len, d_model] memory: 编码器输出记忆[batch_size, src_seq_len, d_model] src_mask: 源序列掩码 tgt_mask: 目标序列掩码填充掩码 前瞻掩码 m memory # 第一子层带掩码的多头自注意力 x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)) # 第二子层编码-解码多头注意力 # Q来自解码器上一层的输出K和V来自编码器的输出memory x self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask)) # 第三子层前馈网络 x self.sublayer[2](x, self.feed_forward) return x解码器的两个关键点self_attn在训练时必须传入tgt_mask它是目标序列的填充掩码和前瞻掩码的逻辑与防止模型看到未来信息。src_attn这是解码器“询问”编码器的环节。查询Q来自解码器自身的输出而键K和值V来自编码器的最终输出memory。src_mask确保了解码器不会关注源序列中的填充位置。5.3 编码器与解码器堆叠最后我们将多个编码器层和解码器层堆叠起来并加上最终的嵌入和线性输出层。class Encoder(nn.Module): def __init__(self, layer, N): super(Encoder, self).__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.size) # 最后再加一个LayerNormPost-LN结构 def forward(self, x, mask): for layer in self.layers: x layer(x, mask) return self.norm(x) class Decoder(nn.Module): def __init__(self, layer, N): super(Decoder, self).__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.size) def forward(self, x, memory, src_mask, tgt_mask): for layer in self.layers: x layer(x, memory, src_mask, tgt_mask) return self.norm(x)注意在原始的Transformer论文Post-LN中每个子层输出后先做AddNorm最后整个编码器/解码器输出后再做一次LayerNorm。有些现代变体如Pre-LN会将LayerNorm放在子层之前这通常能使训练更稳定。6. 完整Transformer模型与输出生成现在我们可以将所有部分组合成最终的Transformer模型。6.1 模型组装class Transformer(nn.Module): def __init__(self, encoder, decoder, src_embed, tgt_embed, generator): super(Transformer, self).__init__() self.encoder encoder self.decoder decoder self.src_embed src_embed # 源语言嵌入 位置编码 self.tgt_embed tgt_embed # 目标语言嵌入 位置编码 self.generator generator # 最后的线性层 softmax def encode(self, src, src_mask): return self.encoder(self.src_embed(src), src_mask) def decode(self, tgt, memory, src_mask, tgt_mask): return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask) def forward(self, src, tgt, src_mask, tgt_mask): # 训练时使用 memory self.encode(src, src_mask) output self.decode(tgt, memory, src_mask, tgt_mask) return self.generator(output) # 映射到词汇表概率generator通常是一个简单的线性层将解码器输出的d_model维向量映射到目标词汇表大小的维度后面接一个LogSoftmax或CrossEntropyLoss内部处理。class Generator(nn.Module): def __init__(self, d_model, vocab_size): super(Generator, self).__init__() self.proj nn.Linear(d_model, vocab_size) def forward(self, x): return torch.log_softmax(self.proj(x), dim-1)6.2 掩码的创建掩码的创建是Transformer实现中容易出错的一环。def subsequent_mask(size): 生成一个前瞻掩码上三角矩阵。 形状: [size, size] attn_shape (1, size, size) subsequent_mask torch.triu(torch.ones(attn_shape), diagonal1).type(torch.uint8) return subsequent_mask 0 # 将下三角包括对角线设为True上三角设为False # 在训练时构建目标掩码 def make_std_mask(tgt, pad): tgt: 目标序列形状 [batch_size, tgt_len] pad: 填充符的索引 返回: 掩码 [batch_size, 1, tgt_len, tgt_len] tgt_mask (tgt ! pad).unsqueeze(-2) # 填充掩码 [batch_size, 1, tgt_len] tgt_mask tgt_mask subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data) return tgt_masksubsequent_mask创建了一个布尔矩阵其中[i, j]为True表示在预测第i个位置时允许关注第j个位置j i。make_std_mask将填充掩码和这个前瞻掩码结合得到最终用于解码器自注意力层的掩码。6.3 推理时的序列生成在推理如机器翻译时我们通常使用自回归autoregressive的方式逐个生成目标词。def greedy_decode(model, src, src_mask, max_len, start_symbol): 贪心解码 memory model.encode(src, src_mask) # 初始化输出序列以起始符开始 ys torch.ones(1, 1).fill_(start_symbol).type_as(src.data) for i in range(max_len-1): # 为当前已生成的序列创建掩码 tgt_mask subsequent_mask(ys.size(1)).type_as(src.data) # 解码 out model.decode(ys, memory, src_mask, tgt_mask) # 取最后一个时间步的输出并通过generator得到概率 prob model.generator(out[:, -1]) # 选择概率最大的词贪心 _, next_word torch.max(prob, dim1) next_word next_word.item() # 将新词拼接到输出序列 ys torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim1) # 如果生成结束符则停止 if next_word EOS_IDX: # EOS_IDX是结束符索引 break return ys这是最简单的贪心搜索。实际应用中会使用束搜索Beam Search来获得更好的结果其核心思想是每一步保留多个概率最高的候选序列而不是只保留一个。7. 训练技巧与实战心得理解了架构要让模型真正工作起来还需要关注训练细节。7.1 学习率调度与优化器选择Transformer论文使用了Adam优化器并配合一个特殊的学习率调度器——预热Warmup后衰减。class NoamOpt: 优化器包装器实现学习率按步数先线性增加再按步数的反平方根衰减。 def __init__(self, model_size, factor, warmup, optimizer): self.optimizer optimizer self._step 0 self.warmup warmup self.factor factor self.model_size model_size self._rate 0 def step(self): self._step 1 rate self.rate() for p in self.optimizer.param_groups: p[lr] rate self._rate rate self.optimizer.step() def rate(self, stepNone): if step is None: step self._step return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))这个调度器在训练早期step warmup线性增加学习率有助于稳定训练之后按step的逆平方根衰减。factor是一个缩放因子model_size即d_model。通常与Adam优化器一起使用beta10.9,beta20.98,epsilon1e-9。7.2 标签平滑与正则化标签平滑Label Smoothing是Transformer中的一个重要技巧。它将硬标签如[0, 0, 1, 0]替换为软标签如[0.1/4, 0.1/4, 0.9, 0.1/4]其中epsilon是平滑因子常取0.1。这可以防止模型对正确答案过于自信起到正则化作用提升泛化能力。在PyTorch中可以直接使用nn.CrossEntropyLoss的label_smoothing参数1.10.0版本。Dropout在Transformer中广泛应用在嵌入层与位置编码相加后、注意力权重计算后、FFN的激活函数后。这是防止过拟合的关键。7.3 梯度裁剪由于Transformer层数可能很深梯度爆炸是一个潜在风险。在每一步优化器更新之前进行梯度裁剪是标准操作。max_grad_norm 1.0 # 梯度范数阈值 torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) optimizer.step()7.4 一个常见的坑张量维度与设备在编写和调试时最常遇到的错误是张量维度不匹配或设备不一致CPU vs GPU。务必注意注意力分数计算时mask的维度需要能广播到scores的维度[batch_size, num_heads, seq_len, seq_len]。在将数据送入模型前确保src,tgt,src_mask,tgt_mask都在同一个设备上model.to(device)后输入数据也要tensor.to(device)。使用torch.autograd.set_detect_anomaly(True)可以在前向传播或反向传播出错时提供更详细的错误追踪信息对调试非常有帮助。从零开始实现一个完整的Transformer是一次深刻的学习体验。它迫使你理解每一个公式如何转化为可运行的代码每一个超参数如何影响模型行为。虽然现在有许多优秀的库如Hugging Face Transformers提供了现成的实现但亲手搭建一遍会让你在面对更复杂的模型变体或需要自定义修改时拥有十足的底气。希望这份详细的代码解读能成为你探索Transformer世界的一块坚实垫脚石。
1. 从理论到实践为什么需要逐行解读Transformer代码如果你已经看过不少关于Transformer的论文解读和架构图但打开PyTorch实现代码时依然感觉一头雾水那么这篇文章就是为你准备的。理论上的多头自注意力、前馈网络、层归一化在代码里是如何组织成一个高效、可训练的模块的位置编码怎么和词嵌入相加训练和推理时的掩码机制有何不同这些细节光看公式和框图是学不到的。我自己在初次实现Transformer时就踩过不少坑比如忘了在解码器的自注意力层应用前瞻掩码导致模型在训练时“作弊”又比如对层归一化的affine参数理解不透影响了模型收敛。通过这篇超详细的代码解读我希望带你穿透理论直抵工程实现的精髓。我们将基于PyTorch从最基础的张量操作开始一步步搭建一个完整的Transformer模型并解释每一行代码背后的设计意图和数学原理。无论你是想深入理解大模型的基础架构还是准备动手复现或修改Transformer用于自己的任务这篇文章都将提供一份可靠的“地图”。2. 环境搭建与项目结构不仅仅是安装PyTorch在深入代码之前一个清晰、可复现的环境是基石。很多人以为环境搭建就是pip install torch实则不然。2.1 PyTorch版本选择与安装避坑首先访问PyTorch官网获取安装命令是最稳妥的方式。你需要根据你的CUDA版本如果使用GPU来选择。例如对于CUDA 12.1你可能需要安装类似torch2.1.0cu121的版本。一个常见的错误是使用conda安装时默认安装的是CPU版本导致后续无法利用GPU加速。注意如果你遇到类似“invalid archive error”的错误这通常是因为网络问题导致安装包下载不完整。解决方法包括使用国内镜像源如清华源、阿里云源或者直接下载whl文件进行离线安装。我建议创建一个独立的Conda环境来管理项目依赖conda create -n transformer_tutorial python3.9 conda activate transformer_tutorial # 根据你的CUDA版本从官网复制安装命令例如 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121除了PyTorch我们还需要一些辅助库pip install numpy matplotlib tqdmmatplotlib用于可视化注意力权重等tqdm可以给训练循环添加进度条提升体验。2.2 代码文件组织一个清晰的项目结构有助于管理复杂度。建议创建如下目录transformer_code/ ├── model.py # Transformer模型定义核心 ├── train.py # 训练脚本 ├── inference.py # 推理/验证脚本 ├── config.py # 超参数配置模型尺寸、层数等 ├── data_loader.py # 数据加载与预处理 ├── utils.py # 工具函数如位置编码可视化 └── main.py # 主程序入口在model.py中我们将按模块构建Transformer嵌入层、位置编码、多头注意力、前馈网络、编码器层、解码器层最后组装成完整的Encoder-Decoder。这种模块化的方式不仅代码清晰也便于单独测试每个组件。3. 核心模块拆解从嵌入层到注意力机制Transformer的威力源于其精巧的模块设计。我们将自底向上看看每个模块在PyTorch中如何实现。3.1 词嵌入与位置编码给模型注入“顺序”信息Transformer本身不具备处理序列顺序的能力因此需要位置编码Positional Encoding。在代码中这通常是两个独立的层最终相加。词嵌入层Embedding非常简单就是PyTorch的nn.Embeddingimport torch.nn as nn class Embeddings(nn.Module): def __init__(self, vocab_size, d_model): super(Embeddings, self).__init__() self.lut nn.Embedding(vocab_size, d_model) self.d_model d_model # 模型维度如512 def forward(self, x): # x: [batch_size, seq_len] # 输出: [batch_size, seq_len, d_model] return self.lut(x) * math.sqrt(self.d_model) # 缩放参见论文这里乘以math.sqrt(self.d_model)是原论文中的一个小技巧目的是在与位置编码相加前保持嵌入向量的方差在一个合适的范围。位置编码层Positional Encoding的实现则更有趣。它使用正弦和余弦函数来生成绝对位置信息import torch import math class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout0.1, max_len5000): super(PositionalEncoding, self).__init__() self.dropout nn.Dropout(pdropout) # 计算位置编码矩阵形状为 [max_len, d_model] pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # [max_len, 1] div_term torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # 对偶数索引应用sin奇数索引应用cos 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] 便于广播 # 将pe注册为buffer不参与训练但会随模型保存/加载 self.register_buffer(pe, pe) def forward(self, x): # x: [batch_size, seq_len, d_model] x x self.pe[:, :x.size(1)] # 只取前seq_len个位置 return self.dropout(x)关键点在于div_term的计算。公式exp(2i * (-log(10000)/d_model))等价于1/(10000^(2i/d_model))这是原论文中使用的频率项。这样每个位置、每个维度都有一个独特的编码。register_buffer确保这个计算好的pe矩阵能被保存和加载但它不是可训练参数。3.2 缩放点积注意力注意力机制的心脏这是整个Transformer最核心的运算单元。其输入是查询Q、键K、值V三个矩阵。def attention(query, key, value, maskNone, dropoutNone): 计算缩放点积注意力。 参数: query, key, value: 形状均为 [batch_size, num_heads, seq_len, d_k] mask: 可选的掩码形状为 [batch_size, 1, seq_len, seq_len] 或 [batch_size, seq_len, seq_len] dropout: nn.Dropout层实例 返回: 加权后的输出注意力权重 d_k query.size(-1) # 获取键向量的维度 # 计算QK^T并缩放 scores torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: # 将mask中为True或1的位置替换为一个极小的负数这样在softmax后权重接近0 scores scores.masked_fill(mask 0, -1e9) p_attn torch.softmax(scores, dim-1) # 在最后一个维度键序列上做softmax if dropout is not None: p_attn dropout(p_attn) # 用注意力权重加权值向量 return torch.matmul(p_attn, value), p_attn为什么需要缩放除以sqrt(d_k)当d_k较大时点积的结果可能会变得非常大这将softmax函数推入梯度极小的区域导致训练困难梯度消失。缩放操作可以稳定梯度。掩码Mask的作用至关重要。在Transformer中主要有两种掩码填充掩码Padding Mask用于忽略序列中无效的填充位置pad。通常在注意力分数计算前将填充位置对应的分数置为一个很大的负数如-1e9。前瞻掩码Look-ahead Mask仅用于解码器的自注意力层。目的是防止解码器在预测第t个词时“看到”第t个词之后的信息即未来信息。这是一个上三角矩阵对角线及以下为1以上为0。3.3 多头注意力并行化的艺术单头注意力可能只关注到序列中特定模式的依赖关系。多头注意力允许模型同时关注来自不同表示子空间的信息。class MultiHeadedAttention(nn.Module): def __init__(self, h, d_model, dropout0.1): h: 头的数量 d_model: 模型总维度 super(MultiHeadedAttention, self).__init__() assert d_model % h 0 # 确保d_model能被h整除 self.d_k d_model // h # 每个头的维度 self.h h # 定义四个线性层用于生成Q, K, V和最终的输出投影 self.linears nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(4)]) self.attn None # 用于保存注意力权重便于可视化或分析 self.dropout nn.Dropout(pdropout) def forward(self, query, key, value, maskNone): if mask is not None: # 同样的掩码需要应用于所有头 mask mask.unsqueeze(1) # [batch_size, 1, 1, seq_len] 或 [batch_size, 1, seq_len, seq_len] batch_size query.size(0) # 1) 线性投影并分头 # 对Q, K, V分别做线性变换然后重塑形状以分离出头 query, key, value [ lin(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2) for lin, x in zip(self.linears, (query, key, value)) ] # 此时 query, key, value 形状: [batch_size, num_heads, seq_len, d_k] # 2) 在分头后的数据上应用注意力函数 x, self.attn attention(query, key, value, maskmask, dropoutself.dropout) # x形状: [batch_size, num_heads, seq_len, d_k] # 3) 合并多头 x x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) # x形状: [batch_size, seq_len, d_model] # 4) 通过最后一个线性层进行输出投影 return self.linears[-1](x)分头与合并的细节代码中view和transpose的操作顺序是关键。先通过线性层将d_model维的输入映射到d_model维然后通过view重塑为[batch_size, seq_len, num_heads, d_k]再通过transpose(1,2)交换seq_len和num_heads维度得到[batch_size, num_heads, seq_len, d_k]这样每个头就能独立计算注意力。合并时先transpose回来再用view合并。4. 前馈网络与残差连接稳定训练的保障注意力层捕捉的是序列元素间的关系而前馈网络Feed-Forward Network, FFN则对每个位置的表示进行独立的、复杂的非线性变换。4.1 位置级前馈网络FFN是一个简单的两层全连接网络中间有一个ReLU激活函数。class PositionwiseFeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 nn.Linear(d_model, d_ff) # 膨胀层通常d_ff4*d_model self.w_2 nn.Linear(d_ff, d_model) # 收缩层 self.dropout nn.Dropout(dropout) def forward(self, x): # x: [batch_size, seq_len, d_model] # 先膨胀再激活再Dropout最后收缩 return self.w_2(self.dropout(torch.relu(self.w_1(x))))为什么需要这个FFN自注意力层本质上是线性变换的加权和尽管权重是动态计算的。加入带有非线性激活函数如ReLU的FFN极大地增强了模型的表示能力使其能够学习更复杂的模式。4.2 层归一化与残差连接训练深度网络的利器这是Transformer能够成功堆叠多层如12层编码器的关键。每个子层自注意力、FFN都被包裹在一个“残差连接层归一化”的模块中。class SublayerConnection(nn.Module): 一个残差连接后接层归一化。 注意为了简化原论文是先做层归一化再进入子层。但后续研究如Pre-LN Transformer发现先归一化效果可能更好。 这里实现原论文的 Post-LN 结构。 def __init__(self, size, dropout): super(SublayerConnection, self).__init__() self.norm nn.LayerNorm(size) self.dropout nn.Dropout(dropout) def forward(self, x, sublayer): x: 输入张量 sublayer: 一个函数它本身定义了一个子层如多头注意力或FFN # 原论文结构x - LayerNorm - Sublayer - Dropout - Add # 即残差连接应用于子层输出和原始输入x之间 return x self.dropout(sublayer(self.norm(x)))层归一化LayerNorm vs 批归一化BatchNorm对于序列数据BatchNorm会沿着批次维度归一化这在不同序列长度和批次内样本差异大时不稳定。LayerNorm则对每个样本、每个时间步的特征维度进行归一化更适合序列模型。nn.LayerNorm默认启用affine参数即可学习的缩放和偏移这很重要因为它让模型能够学习到最适合当前层的分布。残差连接Residual Connection允许梯度直接流过网络极大地缓解了深度网络中的梯度消失问题。公式output x sublayer(norm(x))是标准做法。5. 编码器与解码器组装核心模块有了上面的积木我们现在可以搭建编码器层和解码器层然后将它们堆叠起来。5.1 编码器层自注意力与FFN的堆叠一个编码器层包含一个多头自注意力子层和一个前馈网络子层每个子层都有残差连接和层归一化。class EncoderLayer(nn.Module): def __init__(self, size, self_attn, feed_forward, dropout): super(EncoderLayer, self).__init__() self.self_attn self_attn self.feed_forward feed_forward self.sublayer nn.ModuleList([SublayerConnection(size, dropout) for _ in range(2)]) self.size size # d_model def forward(self, x, mask): x: 编码器输入 [batch_size, src_seq_len, d_model] mask: 源序列的掩码通常是填充掩码[batch_size, 1, src_seq_len] # 第一个子层多头自注意力自注意力意味着 Q, K, V 都来自同一个输入x x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 第二个子层前馈网络 x self.sublayer[1](x, self.feed_forward) return x注意编码器的自注意力中查询Q、键K、值V都来自上一层编码器的输出对于第一层则是输入嵌入与位置编码的和。mask参数用于屏蔽填充符。5.2 解码器层带掩码的自注意力与编码-解码注意力解码器层比编码器层多一个注意力子层即编码-解码注意力层又称交叉注意力层。class DecoderLayer(nn.Module): def __init__(self, size, self_attn, src_attn, feed_forward, dropout): super(DecoderLayer, self).__init__() self.size size self.self_attn self_attn # 带掩码的自注意力 self.src_attn src_attn # 编码-解码注意力 self.feed_forward feed_forward # 三个子层连接 self.sublayer nn.ModuleList([SublayerConnection(size, dropout) for _ in range(3)]) def forward(self, x, memory, src_mask, tgt_mask): x: 解码器输入 [batch_size, tgt_seq_len, d_model] memory: 编码器输出记忆[batch_size, src_seq_len, d_model] src_mask: 源序列掩码 tgt_mask: 目标序列掩码填充掩码 前瞻掩码 m memory # 第一子层带掩码的多头自注意力 x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)) # 第二子层编码-解码多头注意力 # Q来自解码器上一层的输出K和V来自编码器的输出memory x self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask)) # 第三子层前馈网络 x self.sublayer[2](x, self.feed_forward) return x解码器的两个关键点self_attn在训练时必须传入tgt_mask它是目标序列的填充掩码和前瞻掩码的逻辑与防止模型看到未来信息。src_attn这是解码器“询问”编码器的环节。查询Q来自解码器自身的输出而键K和值V来自编码器的最终输出memory。src_mask确保了解码器不会关注源序列中的填充位置。5.3 编码器与解码器堆叠最后我们将多个编码器层和解码器层堆叠起来并加上最终的嵌入和线性输出层。class Encoder(nn.Module): def __init__(self, layer, N): super(Encoder, self).__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.size) # 最后再加一个LayerNormPost-LN结构 def forward(self, x, mask): for layer in self.layers: x layer(x, mask) return self.norm(x) class Decoder(nn.Module): def __init__(self, layer, N): super(Decoder, self).__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.size) def forward(self, x, memory, src_mask, tgt_mask): for layer in self.layers: x layer(x, memory, src_mask, tgt_mask) return self.norm(x)注意在原始的Transformer论文Post-LN中每个子层输出后先做AddNorm最后整个编码器/解码器输出后再做一次LayerNorm。有些现代变体如Pre-LN会将LayerNorm放在子层之前这通常能使训练更稳定。6. 完整Transformer模型与输出生成现在我们可以将所有部分组合成最终的Transformer模型。6.1 模型组装class Transformer(nn.Module): def __init__(self, encoder, decoder, src_embed, tgt_embed, generator): super(Transformer, self).__init__() self.encoder encoder self.decoder decoder self.src_embed src_embed # 源语言嵌入 位置编码 self.tgt_embed tgt_embed # 目标语言嵌入 位置编码 self.generator generator # 最后的线性层 softmax def encode(self, src, src_mask): return self.encoder(self.src_embed(src), src_mask) def decode(self, tgt, memory, src_mask, tgt_mask): return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask) def forward(self, src, tgt, src_mask, tgt_mask): # 训练时使用 memory self.encode(src, src_mask) output self.decode(tgt, memory, src_mask, tgt_mask) return self.generator(output) # 映射到词汇表概率generator通常是一个简单的线性层将解码器输出的d_model维向量映射到目标词汇表大小的维度后面接一个LogSoftmax或CrossEntropyLoss内部处理。class Generator(nn.Module): def __init__(self, d_model, vocab_size): super(Generator, self).__init__() self.proj nn.Linear(d_model, vocab_size) def forward(self, x): return torch.log_softmax(self.proj(x), dim-1)6.2 掩码的创建掩码的创建是Transformer实现中容易出错的一环。def subsequent_mask(size): 生成一个前瞻掩码上三角矩阵。 形状: [size, size] attn_shape (1, size, size) subsequent_mask torch.triu(torch.ones(attn_shape), diagonal1).type(torch.uint8) return subsequent_mask 0 # 将下三角包括对角线设为True上三角设为False # 在训练时构建目标掩码 def make_std_mask(tgt, pad): tgt: 目标序列形状 [batch_size, tgt_len] pad: 填充符的索引 返回: 掩码 [batch_size, 1, tgt_len, tgt_len] tgt_mask (tgt ! pad).unsqueeze(-2) # 填充掩码 [batch_size, 1, tgt_len] tgt_mask tgt_mask subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data) return tgt_masksubsequent_mask创建了一个布尔矩阵其中[i, j]为True表示在预测第i个位置时允许关注第j个位置j i。make_std_mask将填充掩码和这个前瞻掩码结合得到最终用于解码器自注意力层的掩码。6.3 推理时的序列生成在推理如机器翻译时我们通常使用自回归autoregressive的方式逐个生成目标词。def greedy_decode(model, src, src_mask, max_len, start_symbol): 贪心解码 memory model.encode(src, src_mask) # 初始化输出序列以起始符开始 ys torch.ones(1, 1).fill_(start_symbol).type_as(src.data) for i in range(max_len-1): # 为当前已生成的序列创建掩码 tgt_mask subsequent_mask(ys.size(1)).type_as(src.data) # 解码 out model.decode(ys, memory, src_mask, tgt_mask) # 取最后一个时间步的输出并通过generator得到概率 prob model.generator(out[:, -1]) # 选择概率最大的词贪心 _, next_word torch.max(prob, dim1) next_word next_word.item() # 将新词拼接到输出序列 ys torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim1) # 如果生成结束符则停止 if next_word EOS_IDX: # EOS_IDX是结束符索引 break return ys这是最简单的贪心搜索。实际应用中会使用束搜索Beam Search来获得更好的结果其核心思想是每一步保留多个概率最高的候选序列而不是只保留一个。7. 训练技巧与实战心得理解了架构要让模型真正工作起来还需要关注训练细节。7.1 学习率调度与优化器选择Transformer论文使用了Adam优化器并配合一个特殊的学习率调度器——预热Warmup后衰减。class NoamOpt: 优化器包装器实现学习率按步数先线性增加再按步数的反平方根衰减。 def __init__(self, model_size, factor, warmup, optimizer): self.optimizer optimizer self._step 0 self.warmup warmup self.factor factor self.model_size model_size self._rate 0 def step(self): self._step 1 rate self.rate() for p in self.optimizer.param_groups: p[lr] rate self._rate rate self.optimizer.step() def rate(self, stepNone): if step is None: step self._step return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))这个调度器在训练早期step warmup线性增加学习率有助于稳定训练之后按step的逆平方根衰减。factor是一个缩放因子model_size即d_model。通常与Adam优化器一起使用beta10.9,beta20.98,epsilon1e-9。7.2 标签平滑与正则化标签平滑Label Smoothing是Transformer中的一个重要技巧。它将硬标签如[0, 0, 1, 0]替换为软标签如[0.1/4, 0.1/4, 0.9, 0.1/4]其中epsilon是平滑因子常取0.1。这可以防止模型对正确答案过于自信起到正则化作用提升泛化能力。在PyTorch中可以直接使用nn.CrossEntropyLoss的label_smoothing参数1.10.0版本。Dropout在Transformer中广泛应用在嵌入层与位置编码相加后、注意力权重计算后、FFN的激活函数后。这是防止过拟合的关键。7.3 梯度裁剪由于Transformer层数可能很深梯度爆炸是一个潜在风险。在每一步优化器更新之前进行梯度裁剪是标准操作。max_grad_norm 1.0 # 梯度范数阈值 torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) optimizer.step()7.4 一个常见的坑张量维度与设备在编写和调试时最常遇到的错误是张量维度不匹配或设备不一致CPU vs GPU。务必注意注意力分数计算时mask的维度需要能广播到scores的维度[batch_size, num_heads, seq_len, seq_len]。在将数据送入模型前确保src,tgt,src_mask,tgt_mask都在同一个设备上model.to(device)后输入数据也要tensor.to(device)。使用torch.autograd.set_detect_anomaly(True)可以在前向传播或反向传播出错时提供更详细的错误追踪信息对调试非常有帮助。从零开始实现一个完整的Transformer是一次深刻的学习体验。它迫使你理解每一个公式如何转化为可运行的代码每一个超参数如何影响模型行为。虽然现在有许多优秀的库如Hugging Face Transformers提供了现成的实现但亲手搭建一遍会让你在面对更复杂的模型变体或需要自定义修改时拥有十足的底气。希望这份详细的代码解读能成为你探索Transformer世界的一块坚实垫脚石。
)
