从PyTorch源码拆解BERT：双向机制、输入表示与预训练任务实现-北京尧图网络

简介面向自然语言处理学习者与PyTorch实践者的BERT实现代码包对应Google AI 2018年提出的深度双向Transformer预训练模型论文适合希望从源码层面理解BERT结构、在本地复现语言模型训练流程的中级开发者参考。压缩包共33个文件以27个Python脚本为主体覆盖模型定义、数据集处理与训练器模块另含Makefile、requirements.txt、配置文件及README说明整体仅28KB轻量易读代码附带简单注释便于逐行梳理。项目当前仍在开发中代码尚未经过完整验证读者可将其作为学习Transformer编码器与预训练思路的入门素材结合语料准备与快速开始说明动手调试理解自注意力机制在语言模型中的应用方式。已有879人学习关注适合作为论文阅读之外的代码对照参考。1. BERT 的“双向”到底双向在哪一个被误读最多的结论很多人第一次接触 BERT会以为“双向”就是把句子倒过来再跑一遍或者把左右两个 LSTM 拼在一起。真正翻开源码才会发现BERT 的双向不是结构上的双向而是训练目标上的双向它把一句话里 15% 的 token 挖掉让模型同时利用左右两侧的上下文去猜这些被挖掉的位置。这个设计直接绕开了自回归语言模型“只能看左边”的限制也让 Transformer 的并行能力第一次在预训练阶段被彻底释放。搞懂这一点才能理解为什么 BERT 的输入要三段拼接、为什么要有 mask、为什么位置编码是学出来的而不是正余弦。这份 PyTorch 实现把论文里的每个模块拆成了可读的类从 Embedding 到 Transformer 堆叠再到两个预训练头逐层对应原论文的公式是少见的能把“机制”和“代码”对上号的资源。2. 输入表示为什么一条句子要拆成三段 Embedding2.1 Token、Segment、Position 三路相加的动机BERT 的输入不是简单地把词转成 id 就完事。源码里BertEmbeddings这个类做了三件事查词表得到 token embedding、查句子段得到 segment embedding、按位置查 position embedding然后把三者逐元素相加。为什么要加而不是拼接因为拼接会让维度膨胀后续每一层注意力的计算量都跟着涨相加则保持维度不变同时让三种信息在同一个向量空间里叠加。这个选择背后有个隐含假设三种信息是正交的、可以线性叠加的。实际训练中这个假设并不严格成立但实验证明够用这也是深度学习里常见的“工程上先跑通再说”的取舍。class BertEmbeddings(nn.Module): def __init__(self, vocab_size, hidden_size, max_position, type_vocab_size): super().__init__() # 词表嵌入把 token id 映射到 hidden_size 维向量 self.word_embeddings nn.Embedding(vocab_size, hidden_size) # 位置嵌入可学习参数长度固定为 max_position self.position_embeddings nn.Embedding(max_position, hidden_size) # 段嵌入区分句子 A 和句子 B通常只有 0/1 两种 self.token_type_embeddings nn.Embedding(type_vocab_size, hidden_size) self.LayerNorm nn.LayerNorm(hidden_size, eps1e-12) self.dropout nn.Dropout(0.1) def forward(self, input_ids, token_type_ids): seq_len input_ids.size(1) # 位置 id 从 0 到 seq_len-1固定生成不参与梯度 position_ids torch.arange(seq_len, dtypetorch.long, deviceinput_ids.device) position_ids position_ids.unsqueeze(0).expand_as(input_ids) words_emb self.word_embeddings(input_ids) pos_emb self.position_embeddings(position_ids) type_emb self.token_type_embeddings(token_type_ids) # 三路相加后统一做 LayerNorm稳定数值分布 embeddings words_emb pos_emb type_emb return self.dropout(self.LayerNorm(embeddings))这段代码里有个容易被忽略的细节position_ids是用torch.arange现场生成的而不是作为输入传进来。这意味着位置信息完全由序列长度决定和具体是哪个句子无关。好处是实现简单坏处是模型学到的位置是“绝对位置”换一个更长的序列就得重新学。后来的 RoPE 之类方案就是冲着这个痛点去的但在这份实现里绝对位置编码已经够用因为预训练的最大长度是固定的 512。2.2 特殊 token 的布局与 mask 的生成输入序列的布局是[CLS] 句子A [SEP] 句子B [SEP]。[CLS]这个位置的最终隐状态会被拿去做下一句预测任务[SEP]用来分隔两个句子。源码里BertModel的 forward 会同时返回sequence_output和pooled_output后者就是[CLS]位置经过一个线性层加 tanh 的结果。attention mask 的生成逻辑是真实 token 位置为 1padding 位置为 0然后在注意力计算时把 0 的位置填成负无穷softmax 之后这些位置的权重就趋近于零。这个机制保证了 padding 不会污染真实 token 的表示是批处理能跑起来的前提。提示mask 的维度要和注意力分数的维度对齐常见错误是把[batch, seq_len]的 mask 直接加到[batch, heads, seq_len, seq_len]的分数上广播规则不对就会静默出错loss 看起来在降但模型学不到东西。3. Transformer 编码层自注意力的实现细节与数值稳定性3.1 多头注意力为什么要除以 sqrt(d_k)BertSelfAttention里最核心的一行是attention_scores attention_scores / math.sqrt(self.attention_head_size)。这个缩放因子不是拍脑袋定的。点积的方差随维度线性增长维度越大点积结果的绝对值越大softmax 之后分布越尖锐梯度越容易消失。除以sqrt(d_k)相当于把方差拉回 1 附近让 softmax 工作在梯度比较健康的区间。这是 Transformer 能堆深的前提之一也是很多复现翻车的地方——忘了除或者除错维度训练前期 loss 会卡住不动。class BertSelfAttention(nn.Module): def __init__(self, hidden_size, num_heads): super().__init__() self.num_heads num_heads self.head_dim hidden_size // num_heads # 一个线性层同时产出 Q、K、V减少参数和计算 self.query nn.Linear(hidden_size, hidden_size) self.key nn.Linear(hidden_size, hidden_size) self.value nn.Linear(hidden_size, hidden_size) def forward(self, hidden_states, attention_mask): batch, seq_len, _ hidden_states.size() # 拆头[batch, seq, hidden] - [batch, heads, seq, head_dim] q self.query(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k self.key(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v self.value(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 缩放点积防止 softmax 饱和 scores torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim) # mask 位置填极小值softmax 后权重趋零 scores scores attention_mask probs torch.softmax(scores, dim-1) context torch.matmul(probs, v) # 拼回头[batch, heads, seq, head_dim] - [batch, seq, hidden] context context.transpose(1, 2).contiguous().view(batch, seq_len, -1) return contextview加transpose这套组合是 PyTorch 里拆头的标准写法但要注意transpose之后张量在内存里不连续直接view会报错所以代码里用了contiguous()。这个细节在调试时经常被忽略报错信息也不直观属于典型的血泪经验。3.2 残差连接与 LayerNorm 的位置每个子层外面都包了一层残差output LayerNorm(x Sublayer(x))。这个顺序叫 Post-LN是原论文的写法。残差的作用是给梯度一条高速公路让深层网络不至于梯度消失LayerNorm 放在残差之后负责把数值拉回稳定区间。后来的研究指出 Pre-LN先 Norm 再进子层训练更稳定不需要 warmup 也能收敛但这份实现忠于原论文用的是 Post-LN。理解这个顺序的差别对排查“训练到一半 loss 突然爆炸”这类问题很有帮助——Post-LN 在深层堆叠时对学习率更敏感warmup 不是可选项而是必需品。3.3 前馈网络与激活函数的选择BertIntermediate和BertOutput组成前馈网络中间维度是hidden_size * 4激活函数用的是 GELU 而不是 ReLU。GELU 在零点附近是平滑的梯度不会像 ReLU 那样在负半轴直接归零实验上对语言模型更友好。这个选择在 2018 年还算新鲜现在已经成了标配。前馈网络的作用是给每个位置独立做一次非线性变换和注意力负责的“跨位置信息交换”形成互补注意力管横向前馈管纵向。4. 预训练任务MLM 与 NSP 的损失怎么算4.1 Masked Language Model 的 15% 与 80/10/10MLM 的流程是随机选 15% 的 token 位置其中 80% 替换成[MASK]10% 替换成随机 token10% 保持不变。为什么要这么绕如果全换成[MASK]模型只在看到[MASK]时才学习预测而下游微调时根本没有[MASK]会造成预训练和微调的不一致。加入随机替换和保持不变是让模型对“输入可能是噪声”这件事有鲁棒性。源码里BertPreTrainingHeads的predictions是一个hidden_size - vocab_size的线性层损失只在被选中的 15% 位置上计算其余位置忽略。这个“只在部分位置算 loss”的实现是很多人复现时对不齐论文指标的原因——如果把所有位置都算进去loss 数值会小很多但模型学到的信号被稀释了。def mask_tokens(inputs, tokenizer, mlm_prob0.15): labels inputs.clone() # 按 15% 概率选出要预测的位置 prob_matrix torch.full(labels.shape, mlm_prob) masked_indices torch.bernoulli(prob_matrix).bool() # 未被选中的位置 label 设为 -100交叉熵会忽略 labels[~masked_indices] -100 # 80% 换成 [MASK] indices_replaced torch.bernoulli(torch.full(labels.shape, 0.8)).bool() masked_indices inputs[indices_replaced] tokenizer.convert_tokens_to_ids([MASK]) # 10% 换成随机 token indices_random torch.bernoulli(torch.full(labels.shape, 0.5)).bool() masked_indices ~indices_replaced random_words torch.randint(len(tokenizer), labels.shape, dtypetorch.long) inputs[indices_random] random_words[indices_random] # 剩下 10% 保持不变 return inputs, labels注意这里-100是 PyTorchCrossEntropyLoss的默认忽略值不是随便选的数字。参数mlm_prob控制掩码比例论文用的是 0.15实际复现时有人试过 0.2 也能 work但偏离太多会改变任务难度。4.2 Next Sentence Prediction 与它的争议NSP 任务拿[CLS]的输出去判断两个句子是不是连续的。正样本是原文里相邻的两句负样本是从别的文档随机抽的一句。源码里seq_relationship是一个二分类线性层。这个任务后来被不少工作质疑认为它太简单、对下游帮助有限RoBERTa 直接把它去掉了。但在这份实现里保留 NSP 是有价值的它让[CLS]这个位置被迫承载句子级别的语义下游做句子分类时直接拿pooled_output就能用。理解这个设计动机比单纯记住“NSP 有用/没用”的结论更重要。5. 避坑与排查复现 BERT 时最容易翻车的几处5.1 loss 不降或降得极慢现象是训练跑了几个 epochloss 在 10 附近晃悠。原因通常是学习率没配 warmupPost-LN 结构在训练初期梯度方差大直接上大学习率会把参数推到坏区域。解决办法是加线性 warmup前 10% 的步数把学习率从 0 线性拉到峰值再线性衰减。另一个常见原因是 mask 的加法方向搞反了把-10000加到了真实 token 上模型只能看到 padding自然学不到东西。5.2 显存爆掉但 batch size 已经很小现象是CUDA out of memory但 batch 只设了 8。原因往往不是 batch 本身而是注意力分数矩阵的维度是[batch, heads, seq, seq]seq 长度 512、heads 12 时单个样本的这一层就要占几十 MB反向传播还要存一份。解决办法是缩短序列长度、减少 heads或者用梯度累积模拟大 batch。还有一个隐蔽原因是attention_mask用了 float 而不是 boolPyTorch 会额外分配内存做类型转换。5.3 微调时准确率不如随机猜现象是下游任务微调后指标和瞎猜差不多。原因通常是[CLS]的用法错了做句子分类应该取pooled_output已经过线性层和 tanh而不是sequence_output[:, 0, :]。两者差一个线性变换直接取后者等于让分类头自己去学这个变换小数据集上根本学不出来。另一个原因是学习率沿用了预训练的大值微调阶段学习率应该小一到两个数量级否则预训练学到的表示会被冲掉。5.4 位置编码越界现象是换了一个比预训练更长的输入直接报 index out of range。原因是position_embeddings的长度固定为max_position超出就查不到。解决办法是截断输入或者换用可外推的位置编码方案。这个坑在长文本场景下几乎必踩提前知道能省不少调试时间。6. 从这份实现看预训练架构的设计取舍把这份 PyTorch 实现读完会发现 BERT 的每个设计选择背后都有一组权衡。双向注意力换来了更强的上下文表示代价是不能像 GPT 那样做自回归生成绝对位置编码实现简单代价是长度外推能力差Post-LN 忠于原论文代价是对学习率和 warmup 更敏感MLM 的 80/10/10 缓解了预训练与微调的分布差异代价是训练信号只来自 15% 的位置收敛更慢。这些取舍没有绝对的对错只有场景匹配与否。设计点选择代价适用场景注意力方向双向不能直接做生成理解类任务、分类、抽取位置编码可学习绝对位置长度外推差定长输入、短文本Norm 位置Post-LN对学习率敏感忠于原论文的复现预训练目标MLM NSP训练慢、NSP 有争议需要句子级表示验证一份实现是否真的对我一般会做三件事先用极小的模型和造出来的数据过拟合看 loss 能不能降到接近零这一步能排除大部分实现 bug再检查 mask 和 label 的忽略逻辑确认梯度只来自该来的地方最后拿预训练权重跑一个下游小任务看指标是否在合理区间。这三步走完基本能判断这份代码是“看起来像”还是“真的能跑”。从那以后我每次读预训练模型的源码都会先把输入表示和损失计算这两块单独拎出来跑一遍小数据确认数值对得上再往下看。很多玄学问题其实在第一步就能暴露。希望这份拆解能帮你在读这份实现时少走几段弯路把注意力放在机制本身而不是调试细节上。本文还有配套的精品资源点击获取

从PyTorch源码拆解BERT：双向机制、输入表示与预训练任务实现

相关推荐

Rex-Omni：动态特征金字塔与多任务统一的目标检测新范式

YOLOv5双目测距实战：从标定到三维坐标的完整链路

装个浏览器插件它能看见什么：权限清单怎么读

更多资讯

/goal设计：目标驱动的AI编程协作范式

武汉大学数理统计课程核心考点与备考策略

MCP Server：让GPU服务器无缝接入AI编程工具链

agency-agents：智能体系统设计的能动性与协作架构

从命令行到主循环：Hermes引擎源码执行链路全解析

LobeChat本地部署：Docker、环境变量与反向代理实战

今日更新

防幻觉智能体架构：电源硬件设计中的大模型编排与工具化实践

OpenGL面试核心考点：渲染管线、矩阵变换与性能排查全解析

全国计算机二级C语言真题PDF高效刷题指南

本周精选

防幻觉智能体架构：电源硬件设计中的大模型编排与工具化实践

OpenGL面试核心考点：渲染管线、矩阵变换与性能排查全解析

全国计算机二级C语言真题PDF高效刷题指南

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

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