Transformer 原理与实现（二）：从代码看透 Transformer

在上一篇文章 [Transformer 原理与实现一从 Attention 到编码解码机制](https://blog.csdn.net/Cha0DD/article/details/159753362) 中我们从概念层面深入理解了 Transformer 的核心机制。今天我们将通过实际的 PyTorch 代码逐行剖析 Transformer 的实现细节让理论落地真正看透 Transformer 是如何工作的。一、位置编码给词向量加上坐标在上集中我们提到模型无法判断向量顺序所以需要给每个词向量加上一个位置编码。让我们看看如何用代码实现这个坐标系统。1.1 数学公式回顾位置编码使用正弦和余弦函数偶数维度PE(pos, 2i) sin(pos / 10000^(2i/d_model))奇数维度PE(pos, 2i1) cos(pos / 10000^(2i/d_model))1.2 代码实现class PositionalEncoding(nn.Module): 为输入序列的词嵌入向量添加位置编码。 def __init__(self, d_model: int, dropout: float 0.1, max_len: int 5000): super().__init__() self.dropout nn.Dropout(pdropout) # 创建一个足够长的位置编码矩阵 position torch.arange(max_len).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) # pe (positional encoding) 的大小为 (max_len, d_model) pe torch.zeros(max_len, d_model) # 偶数维度使用 sin, 奇数维度使用 cos pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) # 将 pe 注册为 buffer这样它就不会被视为模型参数但会随模型移动 self.register_buffer(pe, pe.unsqueeze(0)) def forward(self, x: torch.Tensor) - torch.Tensor: # x.size(1) 是当前输入的序列长度 # 将位置编码加到输入向量上 x x self.pe[:, :x.size(1)] return self.dropout(x)1.3 代码逐行解析第一步创建位置索引position torch.arange(max_len).unsqueeze(1)这行代码创建了一个从 0 到 max_len-1 的序列position [[0], [1], [2], ... [max_len-1]]第二步计算除数项div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))这行代码计算了每个维度的除数项torch.arange(0, d_model, 2)生成[0, 2, 4, ..., d_model-2]* (-math.log(10000.0) / d_model)乘以缩放因子torch.exp()取指数得到10000^(0/d_model), 10000^(2/d_model), ...第三步填充位置编码pe[:, 0::2] torch.sin(position * div_term) # 偶数维度 pe[:, 1::2] torch.cos(position * div_term) # 奇数维度这里使用了切片操作[:, 0::2]选择所有行的偶数列[:, 1::2]选择所有行的奇数列第四步使用 register_bufferself.register_buffer(pe, pe.unsqueeze(0))这是一个重要的 PyTorch 技巧register_buffer()注册的不是模型参数不会被梯度更新但会随模型一起移动如model.to(device)保存模型时会一起保存加载时自动恢复第五步前向传播def forward(self, x: torch.Tensor) - torch.Tensor: x x self.pe[:, :x.size(1)] # 只取前 seq_len 个位置编码 return self.dropout(x)关键点self.pe[:, :x.size(1)]动态截取适应不同长度的输入直接相加因为词嵌入和位置编码维度相同最后应用 dropout 进行正则化1.4 运行示例# 创建位置编码层 pos_encoder PositionalEncoding(d_model512, dropout0.1) # 模拟输入batch_size2, seq_len10, d_model512 x torch.randn(2, 10, 512) # 应用位置编码 output pos_encoder(x) print(f输入形状: {x.shape}) # torch.Size([2, 10, 512]) print(f输出形状: {output.shape}) # torch.Size([2, 10, 512])二、多头注意力让模型多角度观察多头注意力是 Transformer 的核心创新让我们通过代码理解它是如何工作的。2.1 代码结构class MultiHeadAttention(nn.Module): 多头注意力机制模块 def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads # 每个头的维度 # 定义 Q, K, V 和输出的线性变换层 self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model)2.2 关键点解析维度检查assert d_model % num_heads 0这确保 d_model 可以被 num_heads 整除例如d_model512, num_heads8 → 每个头 64 维d_model768, num_heads12 → 每个头 64 维线性变换层self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model)每个线性层学习一个权重矩阵将输入从 d_model 维映射到 d_model 维。2.3 缩放点积注意力def scaled_dot_product_attention(self, Q, K, V, maskNone): # 1. 计算注意力得分 (QK^T) attn_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # 2. 应用掩码 (如果提供) if mask is not None: attn_scores attn_scores.masked_fill(mask 0, -1e9) # 3. 计算注意力权重 (Softmax) attn_probs torch.softmax(attn_scores, dim-1) # 4. 加权求和 (权重 * V) output torch.matmul(attn_probs, V) return output第一步计算注意力分数attn_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)Q和K的形状都是(batch_size, num_heads, seq_len, d_k)K.transpose(-2, -1)交换最后两个维度变成(batch_size, num_heads, d_k, seq_len)矩阵相乘得到(batch_size, num_heads, seq_len, seq_len)的注意力分数除以sqrt(d_k)进行缩放防止数值过大第二步应用掩码if mask is not None: attn_scores attn_scores.masked_fill(mask 0, -1e9)掩码用于屏蔽某些位置如解码器中的未来信息masked_fill(mask 0, -1e9)将 mask 为 0 的位置设为 -1e9-1e9 是一个非常小的负数经过 softmax 后接近 0第三步Softmax 归一化attn_probs torch.softmax(attn_scores, dim-1)在最后一个维度seq_len上做 softmax结果是每个位置对其他位置的注意力权重总和为 1第四步加权求和output torch.matmul(attn_probs, V)注意力权重(batch_size, num_heads, seq_len, seq_len)V(batch_size, num_heads, seq_len, d_k)输出(batch_size, num_heads, seq_len, d_k)2.4 分头和合头def split_heads(self, x): # 将输入 x 的形状从 (batch_size, seq_length, d_model) # 变换为 (batch_size, num_heads, seq_length, d_k) batch_size, seq_length, d_model x.size() return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2) def combine_heads(self, x): 将输入 x 的形状从 (batch_size, num_heads, seq_length, d_k) 变回 (batch_size, seq_length, d_model) batch_size, num_heads, seq_length, d_k x.size() return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)split_heads 流程输入: (batch_size, seq_len, d_model) (2, 10, 512) view: (2, 10, 8, 64) # 重塑为 4D 张量 transpose(1, 2): (2, 8, 10, 64) # 交换维度让 num_heads 在第二个维度combine_heads 流程输入: (batch_size, num_heads, seq_len, d_k) (2, 8, 10, 64) transpose(1, 2): (2, 10, 8, 64) # 交换回来 contiguous(): 确保内存连续 # transpose 后内存不连续需要连续化 view: (2, 10, 512) # 重塑回原始形状2.5 前向传播def forward(self, Q, K, V, maskNone): # 1. 对 Q, K, V 进行线性变换 Q self.split_heads(self.W_q(Q)) K self.split_heads(self.W_k(K)) V self.split_heads(self.W_v(V)) # 2. 计算缩放点积注意力 attn_output self.scaled_dot_product_attention(Q, K, V, mask) # 3. 合并多头输出并进行最终的线性变换 output self.W_o(self.combine_heads(attn_output)) return output完整流程输入 Q, K, V (每个: batch_size × seq_len × d_model) ↓ 线性变换 W_q, W_k, W_v ↓ split_heads: (batch_size, num_heads, seq_len, d_k) ↓ scaled_dot_product_attention (每个头独立计算) ↓ combine_heads: (batch_size, seq_len, d_model) ↓ 线性变换 W_o ↓ 输出 (batch_size, seq_len, d_model)三、前馈网络把直线掰弯在 上一集中我们提到前馈网络FFN负责非线性变换让模型能够拟合复杂的语言规律。3.1 代码实现class PositionWiseFeedForward(nn.Module): 位置前馈网络模块 def __init__(self, d_model, d_ff, dropout0.1): super(PositionWiseFeedForward, self).__init__() self.linear1 nn.Linear(d_model, d_ff) self.dropout nn.Dropout(dropout) self.linear2 nn.Linear(d_ff, d_model) self.relu nn.ReLU() def forward(self, x): # x 形状: (batch_size, seq_len, d_model) x self.linear1(x) x self.relu(x) x self.dropout(x) x self.linear2(x) # 最终输出形状: (batch_size, seq_len, d_model) return x3.2 逐层解析第一层升维x self.linear1(x)输入:(batch_size, seq_len, d_model)输出:(batch_size, seq_len, d_ff)d_ff 通常设置为 d_model 的 4 倍如 2048升维的目的提供更大的表示空间让模型能够学习更复杂的特征。激活函数x self.relu(x)ReLU (Rectified Linear Unit):max(0, x)引入非线性让网络能够拟合任意函数根据万能近似定理只要有非线性激活神经网络可以逼近任意连续函数Dropoutx self.dropout(x)随机将一部分神经元输出设为 0防止过拟合提高泛化能力只在训练时生效推理时自动关闭第二层降维x self.linear2(x)输入:(batch_size, seq_len, d_ff)输出:(batch_size, seq_len, d_model)恢复到原始维度3.3 为什么需要 FFN注意力机制虽然强大但本质上仍然是线性变换。FFN 通过非线性激活让模型能够特征重组将注意力输出的关系进行重新组合语义提炼从多个关系中提炼出更高级的语义非线性拟合拟合复杂的语言规律无 FFN: 只能画直线线性关系 有 FFN: 可以画曲线非线性关系四、编码器层理解输入语义编码器层结合了多头注意力和前馈网络通过残差连接和层归一化保证训练稳定。4.1 代码实现class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super(EncoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionWiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, mask): # 1. 多头自注意力 attn_output self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_output)) # 2. 前馈网络 ff_output self.feed_forward(x) x self.norm2(x self.dropout(ff_output)) return x4.2 前向传播流程输入 x (batch_size, seq_len, d_model) ↓ 多头自注意力 self_attn(x, x, x, mask) ↓ attn_output ↓ 残差连接: x attn_output ↓ 层归一化 norm1 ↓ 前馈网络 feed_forward ↓ ff_output ↓ 残差连接: x ff_output ↓ 层归一化 norm2 ↓ 输出 (batch_size, seq_len, d_model)4.3 残差连接的作用x self.norm1(x self.dropout(attn_output))残差连接的公式output x F(x)为什么要加残差连接解决梯度消失梯度可以直接流向浅层保留原始信息x 保留原始理解F(x) 学习新关系允许恒等映射F(x) 可以学习到 0相当于该层不生效层归一化LayerNormself.norm1 nn.LayerNorm(d_model)对每个样本的所有特征进行归一化稳定训练过程加速收敛防止数值溢出五、解码器层生成输出解码器层比编码器层多了一个交叉注意力机制用于融合编码器的语义信息。5.1 代码实现class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super(DecoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.cross_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionWiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, encoder_output, src_mask, tgt_mask): # 1. 掩码多头自注意力 (对自己) attn_output self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout(attn_output)) # 2. 交叉注意力 (对编码器输出) cross_attn_output self.cross_attn(x, encoder_output, encoder_output, src_mask) x self.norm2(x self.dropout(cross_attn_output)) # 3. 前馈网络 ff_output self.feed_forward(x) x self.norm3(x self.dropout(ff_output)) return x5.2 三阶段解析阶段 1掩码自注意力attn_output self.self_attn(x, x, x, tgt_mask)Q, K, V 都来自解码器的输入 x使用tgt_mask屏蔽未来信息只能看到当前 token 及其之前的历史信息阶段 2交叉注意力cross_attn_output self.cross_attn(x, encoder_output, encoder_output, src_mask)Q 来自解码器输入 xK, V 来自编码器输出 encoder_output目的是将编码器的语义信息融合到解码器中实现输入输出的语义对齐阶段 3前馈网络ff_output self.feed_forward(x)与编码器层的 FFN 相同进行非线性变换和特征重组5.3 完整流程解码器输入 x (batch_size, tgt_seq_len, d_model) 编码器输出 encoder_output (batch_size, src_seq_len, d_model) ↓ 掩码自注意力: self_attn(x, x, x, tgt_mask) ↓ 残差 归一化 ↓ 交叉注意力: cross_attn(x, encoder_output, encoder_output, src_mask) ↓ 残差 归一化 ↓ 前馈网络 ↓ 残差 归一化 ↓ 输出 (batch_size, tgt_seq_len, d_model)六、完整运行示例让我们运行一个完整的示例看看整个流程是如何工作的。if __name__ __main__: # 设置超参数 d_model 512 num_heads 8 d_ff 2048 dropout 0.1 # 实例化一个编码器层 layer EncoderLayer(d_model, num_heads, d_ff, dropout) # 模拟输入数据: (batch_size2, seq_len10, d_model512) x torch.randn(2, 10, d_model) mask None # 运行测试 try: output layer(x, mask) print(测试成功) print(f输入形状: {x.shape}) print(f输出形状: {output.shape}) except Exception as e: print(f运行出错: {e})6.1 输出结果测试成功 输入形状: torch.Size([2, 10, 512]) 输出形状: torch.Size([2, 10, 512])参数解释d_model 512: 模型的维度num_heads 8: 多头注意力的头数d_ff 2048: 前馈网络的隐藏层维度通常是 d_model 的 4 倍dropout 0.1: Dropout 概率6.2 扩展测试让我们测试更多组件# 测试位置编码 pos_encoder PositionalEncoding(d_model512) x_pos torch.randn(2, 10, 512) output_pos pos_encoder(x_pos) print(f位置编码输出形状: {output_pos.shape}) # 测试多头注意力 mha MultiHeadAttention(d_model512, num_heads8) q torch.randn(2, 10, 512) k torch.randn(2, 10, 512) v torch.randn(2, 10, 512) output_mha mha(q, k, v) print(f多头注意力输出形状: {output_mha.shape}) # 测试前馈网络 ffn PositionWiseFeedForward(d_model512, d_ff2048) x_ffn torch.randn(2, 10, 512) output_ffn ffn(x_ffn) print(f前馈网络输出形状: {output_ffn.shape}) # 测试解码器层 decoder_layer DecoderLayer(d_model512, num_heads8, d_ff2048, dropout0.1) x_decoder torch.randn(2, 5, 512) # 解码器输入 encoder_output torch.randn(2, 10, 512) # 编码器输出 output_decoder decoder_layer(x_decoder, encoder_output, None, None) print(f解码器层输出形状: {output_decoder.shape})6.3 测试输出位置编码输出形状: torch.Size([2, 10, 512]) 多头注意力输出形状: torch.Size([2, 10, 512]) 前馈网络输出形状: torch.Size([2, 10, 512]) 解码器层输出形状: torch.Size([2, 5, 512])七、代码与理论的对应让我们总结一下代码实现与上集中的理论概念的对应关系。7.1 核心组件映射理论概念代码类核心功能位置编码PositionalEncoding给词向量添加位置信息多头注意力MultiHeadAttention让模型多角度观察序列前馈网络PositionWiseFeedForward非线性变换拟合复杂语义编码器层EncoderLayer理解输入语义解码器层DecoderLayer生成输出序列7.2 关键操作映射操作代码作用QKV 变换self.W_q(x)生成查询、键、值向量点积注意力torch.matmul(Q, K.transpose(-2, -1))计算注意力分数缩放/ math.sqrt(self.d_k)防止数值过大Softmaxtorch.softmax(attn_scores, dim-1)归一化为概率加权求和torch.matmul(attn_probs, V)融合信息残差连接x attn_output保留原始信息解决梯度消失层归一化nn.LayerNorm(d_model)稳定训练加速收敛八、本篇总结通过代码逐行分析我们深入理解了 Transformer 的实现细节8.1 关键发现位置编码的巧妙设计使用 sin/cos 函数无需训练但能提供丰富的位置信息多头注意力的分头机制通过split_heads和combine_heads实现高效的多头计算残差连接的重要性x F(x)简单但强大解决了深层网络的梯度消失问题FFN 的非线性能力通过升维 ReLU 降维让模型能够拟合复杂的语言规律8.2 代码优化技巧register_buffer用于存储不需要梯度的张量如位置编码contiguous()在张量转置后使用确保内存连续masked_fill高效的掩码操作用于屏蔽未来信息附完整全代码import torch import torch.nn as nn import math # --- 占位符模块将在后续小节中实现 --- class PositionalEncoding(nn.Module): 为输入序列的词嵌入向量添加位置编码。 def __init__(self, d_model: int, dropout: float 0.1, max_len: int 5000): super().__init__() self.dropout nn.Dropout(pdropout) # 创建一个足够长的位置编码矩阵 position torch.arange(max_len).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) # pe (positional encoding) 的大小为 (max_len, d_model) pe torch.zeros(max_len, d_model) # 偶数维度使用 sin, 奇数维度使用 cos pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) # 将 pe 注册为 buffer这样它就不会被视为模型参数但会随模型移动例如 to(device) self.register_buffer(pe, pe.unsqueeze(0)) def forward(self, x: torch.Tensor) - torch.Tensor: # x.size(1) 是当前输入的序列长度 # 将位置编码加到输入向量上 x x self.pe[:, :x.size(1)] return self.dropout(x) class MultiHeadAttention(nn.Module): 多头注意力机制模块 def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads # 定义 Q, K, V 和输出的线性变换层 self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) def scaled_dot_product_attention(self, Q, K, V, maskNone): # 1. 计算注意力得分 (QK^T) attn_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # 2. 应用掩码 (如果提供) if mask is not None: # 将掩码中为 0 的位置设置为一个非常小的负数这样 softmax 后会接近 0 attn_scores attn_scores.masked_fill(mask 0, -1e9) # 3. 计算注意力权重 (Softmax) attn_probs torch.softmax(attn_scores, dim-1) # 4. 加权求和 (权重 * V) output torch.matmul(attn_probs, V) return output def split_heads(self, x): # 将输入 x 的形状从 (batch_size, seq_length, d_model) # 变换为 (batch_size, num_heads, seq_length, d_k) batch_size, seq_length, d_model x.size() return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2) def combine_heads(self, x): # 将输入 x 的形状从 (batch_size, num_heads, seq_length, d_k) # 变回 (batch_size, seq_length, d_model) batch_size, num_heads, seq_length, d_k x.size() return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model) def forward(self, Q, K, V, maskNone): # 1. 对 Q, K, V 进行线性变换 Q self.split_heads(self.W_q(Q)) K self.split_heads(self.W_k(K)) V self.split_heads(self.W_v(V)) # 2. 计算缩放点积注意力 attn_output self.scaled_dot_product_attention(Q, K, V, mask) # 3. 合并多头输出并进行最终的线性变换 output self.W_o(self.combine_heads(attn_output)) return output class PositionWiseFeedForward(nn.Module): 位置前馈网络模块 def __init__(self, d_model, d_ff, dropout0.1): super(PositionWiseFeedForward, self).__init__() self.linear1 nn.Linear(d_model, d_ff) self.dropout nn.Dropout(dropout) self.linear2 nn.Linear(d_ff, d_model) self.relu nn.ReLU() def forward(self, x): # x 形状: (batch_size, seq_len, d_model) x self.linear1(x) x self.relu(x) x self.dropout(x) x self.linear2(x) # 最终输出形状: (batch_size, seq_len, d_model) return x # --- 编码器核心层 --- class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super(EncoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionWiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, mask): # 残差连接与层归一化将在 3.1.2.4 节中详细解释 # 1. 多头自注意力 attn_output self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_output)) # 2. 前馈网络 ff_output self.feed_forward(x) x self.norm2(x self.dropout(ff_output)) return x # --- 解码器核心层 --- class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super(DecoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.cross_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionWiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, encoder_output, src_mask, tgt_mask): # 1. 掩码多头自注意力 (对自己) attn_output self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout(attn_output)) # 2. 交叉注意力 (对编码器输出) cross_attn_output self.cross_attn(x, encoder_output, encoder_output, src_mask) x self.norm2(x self.dropout(cross_attn_output)) # 3. 前馈网络 ff_output self.feed_forward(x) x self.norm3(x self.dropout(ff_output)) return x if __name__ __main__: # 设置超参数 d_model 512 num_heads 8 d_ff 2048 dropout 0.1 # 实例化一个编码器层 layer EncoderLayer(d_model, num_heads, d_ff, dropout) # 模拟输入数据: (batch_size2, seq_len10, d_model512) x torch.randn(2, 10, d_model) mask None # 运行测试 try: output layer(x, mask) print(测试成功) print(f输入形状: {x.shape}) print(f输出形状: {output.shape}) except Exception as e: print(f运行出错: {e})