简介面向自然语言处理学习者与 PyTorch 实践者的 BERT 复现资源对应 Google AI 2018 年提出的深度双向 Transformer 预训练模型帮助读者在本地理解并跑通语言模型训练流程适合具备一定深度学习基础、希望从代码层面吃透 BERT 结构的中级开发者。压缩包共 33 个文件以 27 个 Python 脚本为主体覆盖模型、数据集与训练器等模块另含 Makefile、requirements.txt、LICENSE、README 及 CI 配置等辅助文件整体约 28KB体量轻便、便于通读。已有 879 人学习下载说明其在入门复现场景中具备一定参考价值。代码附带简单注释结构清晰读者可据此梳理 Transformer 编码器、预训练语料组织与训练入口的实现脉络并借助安装与快速开始说明完成环境搭建和语料准备为后续迁移到具体 NLP 任务打下基础。1. 从一份 BERT 的 PyTorch 复现说起它到底能解决什么问题如果你手头有一批中文工单、合同条款或者客服对话想做一个能真正跑起来的语义分类或实体抽取模型大概率绕不开 BERT 这个底座。而当你打开官方仓库看到的往往是 TensorFlow 时代的 checkpoint 转换脚本和一堆依赖版本说明直接拿来在 PyTorch 里训练并不顺手。这份 BERT-pytorch 实现就是冲着这个落差来的它把 Google 2018 年那版 BERT 的预训练与微调流程用 PyTorch 重新组织成一套可读、可改、可复现的代码覆盖了从词表构建、数据管道、Transformer 编码器堆叠到 MLM/NSP 双任务预训练和下游分类微调的完整链路。适合谁适合已经懂一点 Transformer、但被官方实现劝退过的 NLP 工程师也适合想拿一份干净代码去改结构、做消融实验的研究型开发者。它不是开箱即用的 SaaS而是一份能让你看清每个张量怎么流动的工程底稿。2. 拆开代码包Transformer 编码器与预训练任务的落地结构2.1 目录与模块划分先认清哪块负责什么拿到这类复现包第一件事不是急着跑训练而是把模块边界摸清楚。常见做法是按职责切成四层数据层词表、Dataset、动态 padding 的 collate、模型层Embedding、多头注意力、前馈网络、编码器堆叠、任务头MLM 头、NSP 头、下游分类头、训练层优化器、学习率 warmup、checkpoint 存取。这份实现基本遵循这个划分模型主体集中在一个model目录下预训练入口和微调入口分开配置文件用 JSON 或 YAML 描述超参。这样切的好处是你想换掉 NSP 任务、只保留 MLM改动面被限制在任务头和数据管道两处不用动编码器。我一般会先看三个文件定义BertModel的那个、定义BertForPreTraining的那个、以及train.py里构造 optimizer 和 scheduler 的那段。前两个决定模型能不能对上论文第三个决定训练会不会在 warmup 阶段就发散。下面是一个典型的编码器层结构参数名按常见实现习惯写你对照自己包里的命名微调即可import torch import torch.nn as nn class TransformerBlock(nn.Module): def __init__(self, hidden768, heads12, ff_dim3072, dropout0.1): super().__init__() # 多头自注意力hidden 必须能被 heads 整除 self.attn nn.MultiheadAttention(hidden, heads, dropoutdropout, batch_firstTrue) self.norm1 nn.LayerNorm(hidden) self.norm2 nn.LayerNorm(hidden) # 前馈网络中间维度通常是 hidden 的 4 倍 self.ff nn.Sequential( nn.Linear(hidden, ff_dim), nn.GELU(), # BERT 原文用 GELU不是 ReLU nn.Linear(ff_dim, hidden), ) self.dropout nn.Dropout(dropout) def forward(self, x, attn_maskNone): # Pre-LN 还是 Post-LN 要和权重来源保持一致这里演示 Post-LN attn_out, _ self.attn(x, x, x, attn_maskattn_mask, need_weightsFalse) x self.norm1(x self.dropout(attn_out)) ff_out self.ff(x) x self.norm2(x self.dropout(ff_out)) return x这段代码的关键不在写法而在三个参数hidden决定整个模型的容量base 版是 768large 版是 1024heads是注意力头数base 版 12 头、每头 64 维ff_dim是前馈中间层固定为 hidden 的 4 倍。GELU这个激活函数是 BERT 和早期 Transformer 用 ReLU 的一个明显区别如果你把它换成 ReLU加载预训练权重后效果会掉这是血泪经验。另外batch_firstTrue这个参数要留意PyTorch 原生MultiheadAttention默认是序列维在前很多复现包在这里做过转置处理混用会导致维度对不上。2.2 词表与数据管道WordPiece 和动态 padding 怎么接BERT 用的是 WordPiece 分词词表大小中文常见 21128英文 base 是 30522。这份实现通常自带一个vocab.txt或者一个构建脚本你需要确认词表里的特殊 token 顺序[PAD]、[UNK]、[CLS]、[SEP]、[MASK]的 id 必须和模型 Embedding 的初始化对齐否则预训练阶段 mask 的位置全是错的。数据管道这块预训练需要把一段文本切成句子对按 0.5 的概率决定第二个句子是不是真实的下一句这就是 NSP 标签的来源。from torch.utils.data import Dataset import random class PretrainDataset(Dataset): def __init__(self, docs, tokenizer, max_len128): self.docs docs self.tok tokenizer self.max_len max_len def __len__(self): return len(self.docs) def __getitem__(self, idx): doc self.docs[idx] # 简化按句号切句真实场景用分句器 sents [s for s in doc.split(。) if s.strip()] if len(sents) 2: return self.__getitem__((idx 1) % len(self.docs)) i random.randint(0, len(sents) - 2) sent_a sents[i] # 50% 概率取真实下一句50% 随机取一句 if random.random() 0.5: sent_b, is_next sents[i 1], 1 else: sent_b, is_next random.choice(sents), 0 enc self.tok(sent_a, sent_b, max_lengthself.max_len, truncationTrue, paddingmax_length, return_tensorsNone) input_ids enc[input_ids] # MLM15% 的 token 被选中其中 80% 换 [MASK]10% 换随机词10% 不变 labels [-100] * len(input_ids) for pos, tid in enumerate(input_ids): if tid in (self.tok.cls_token_id, self.tok.sep_token_id, self.tok.pad_token_id): continue if random.random() 0.15: labels[pos] tid r random.random() if r 0.8: input_ids[pos] self.tok.mask_token_id elif r 0.9: input_ids[pos] random.randint(0, self.tok.vocab_size - 1) return {input_ids: input_ids, labels: labels, next_label: is_next}这段逻辑里有两个容易翻车的点。第一labels用-100填充是因为 PyTorch 的CrossEntropyLoss默认ignore_index-100如果你用别的值loss 会被大量无效位置拉偏。第二MLM 的 80/10/10 替换策略必须严格按论文来很多复现为了省事直接全换[MASK]预训练出来的模型在微调时会有明显的预训练-微调不一致问题因为微调阶段根本没有[MASK]。动态 padding 这块训练时按 batch 内最长序列 padding 能省显存但要注意 attention mask 同步生成否则 padding 位置会参与注意力计算。2.3 预训练与微调的衔接权重怎么迁过去预训练跑完你拿到的是一个带 MLM 头和 NSP 头的模型。做下游分类时标准做法是丢掉这两个任务头保留编码器接一个Linear(hidden, num_labels)用[CLS]位置的输出做分类。这里有个实操细节[CLS]的输出在预训练阶段是被 NSP 任务训练过的所以它天然带一点句子级语义这也是 BERT 用[CLS]做分类的合理性来源。迁移时不要重新初始化编码器只初始化分类头学习率上编码器用小学习率常见 2e-5 到 5e-5分类头可以用大一点1e-3 量级这就是常说的分层学习率。from transformers import BertForSequenceClassification, AdamW, get_linear_schedule_with_warmup model BertForSequenceClassification.from_pretrained(./pretrained_ckpt, num_labels8) # 分层学习率编码器小、分类头大 no_decay [bias, LayerNorm.weight] optimizer_grouped [ {params: [p for n, p in model.bert.named_parameters() if not any(nd in n for nd in no_decay)], lr: 3e-5, weight_decay: 0.01}, {params: [p for n, p in model.bert.named_parameters() if any(nd in n for nd in no_decay)], lr: 3e-5, weight_decay: 0.0}, {params: model.classifier.parameters(), lr: 1e-3, weight_decay: 0.01}, ] optimizer AdamW(optimizer_grouped) scheduler get_linear_schedule_with_warmup(optimizer, num_warmup_steps100, num_training_steps2000)weight_decay对 bias 和 LayerNorm 参数置零是常规操作因为这些参数做衰减没有意义还会拖慢收敛。warmup 步数一般取总步数的 10% 左右太少容易在初期震荡太多则浪费训练预算。如果你是从头预训练而不是加载现成权重那 warmup 要更保守常见做法是前 10000 步线性升温。3. 训练参数怎么设从 batch size 到学习率的取舍3.1 显存与 batch size 的换算关系BERT base 在 fp32 下序列长度 128、batch size 32 大概要 10GB 出头显存序列长度 512 时batch size 8 就可能逼近 16GB。这不是玄学是注意力矩阵的平方复杂度决定的。实操里我一般先固定序列长度从 batch size 8 往上试OOM 就减半同时用梯度累积把等效 batch size 补回来。比如目标等效 batch 是 64实际只能跑 8那就累积 8 步再optimizer.step()。注意梯度累积时 loss 要除以累积步数否则梯度会被放大。accum_steps 8 for step, batch in enumerate(loader): outputs model(**batch) loss outputs.loss / accum_steps # 关键归一化 loss.backward() if (step 1) % accum_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad()clip_grad_norm_这个梯度裁剪在 BERT 训练里几乎是标配阈值 1.0 是论文里的常用值。不加的话预训练阶段偶尔会出现 loss 突然飙到 nan尤其是学习率偏大的时候。3.2 学习率、序列长度与训练轮次的组合微调阶段序列长度 128 的分类任务3 到 5 个 epoch 通常就收敛了再多容易过拟合尤其样本量小于一万的时候。学习率 2e-5 到 5e-5 是安全区间超过 1e-4 很容易把预训练权重冲坏。预训练阶段则完全不同学习率常见 1e-4 量级配合 warmup 和线性衰减训练步数动辄几十万步个人设备基本跑不动全量所以这份实现更现实的价值是让你在小语料上做继续预训练continue pretraining而不是从零训一个 base。场景序列长度batch size学习率epoch/步数下游分类微调128323e-53~5 epoch下游分类微调长文本51282e-53~4 epoch小语料继续预训练128165e-51~2 万步从零预训练参考1282561e-4数十万步这张表里的数字不是死的但量级关系要守住序列越长学习率和 batch 都要往下压从零预训练和微调是两个完全不同的预算级别别拿微调的经验去套预训练。3.3 评估指标与早停策略分类任务看 accuracy 和 macro-F1样本不均衡时以 macro-F1 为准。预训练阶段看 MLM 的 loss 和 mask 位置的预测准确率NSP 看二分类准确率但这两个指标只能说明模型在学不能说明它学得好真正的验证还是下游任务。早停我一般按验证集 F1 连续 2 个 epoch 不涨就停同时保存 F1 最高的那个 checkpoint而不是最后一个。这份实现里如果 checkpoint 保存逻辑只存 last建议自己补一个 best 保存不然训练完回头找最优权重会很痛苦。4. 避坑与排查五个真实会卡住你的问题4.1 加载权重后 loss 不降甚至比随机初始化还差现象是微调第一个 epoch 的 loss 在 2.3 附近横盘准确率贴着随机猜。原因多半是权重加载时参数名对不上load_state_dict默认strictTrue会报错但有些人为了跑通改成strictFalse结果编码器一半参数是随机初始化的表面不报错实际模型是残的。解决办法是先打印model.state_dict().keys()和 checkpoint 的 keys 做集合差确认缺失和多余的参数名再决定是改映射还是改模型定义。另一个可能是词表 id 错位[CLS]的 id 在词表里是 101如果你的 tokenizer 配置里写成了别的值[CLS]位置拿到的就是普通词向量。4.2 训练到一半 loss 变 nan现象是前几百步正常突然 loss 变成 nan 且不再恢复。原因通常是学习率过大导致梯度爆炸或者 fp16 混合精度下 loss scale 溢出。解决顺序是先加梯度裁剪阈值 1.0再把学习率降一半如果用了 fp16 就检查GradScaler是否正常unscale_后再裁剪。还有一个隐蔽原因是输入里有全 padding 的样本attention 全被 mask 掉后 softmax 出现 0/0排查时在 collate 里断言每个样本至少有一个有效 token。4.3 显存够但训练速度异常慢现象是 GPU 利用率长期低于 30%一个 epoch 跑得比预期慢好几倍。原因常见于数据加载成了瓶颈DataLoader的num_workers设成 0所有分词和 padding 都在主进程做GPU 一直在等数据。解决办法是把num_workers设成 4 到 8pin_memoryTrue并且把能提前做的分词尽量放到 Dataset 初始化阶段做缓存而不是__getitem__里每次重算。另外动态 padding 如果实现成每个样本单独 padding 到 max_len也会白白浪费算力。4.4 微调效果不如直接拿现成预训练模型现象是自己预训练或继续预训练的模型在下游任务上反而不如直接加载公开权重的效果好。原因往往是小语料继续预训练把原有知识冲掉了也就是灾难性遗忘。解决办法是降低继续预训练的学习率比如 1e-5减少步数或者在 loss 里混入一部分原始语料的 MLM 任务做正则。还有一种情况是预训练语料和下游领域差太远这时候继续预训练的意义本来就不大不如直接微调。4.5 多卡训练后单卡推理结果不一致现象是单卡训练正常换成多卡 DDP 后保存的权重在单卡加载时输出对不上。原因是 BatchNorm 类层在多卡下的统计量同步问题不过 BERT 用的是 LayerNorm理论上不受影响真正的问题多半出在保存时只存了model.module.state_dict()而加载时按单卡结构读key 前缀多了module.。解决办法是保存时统一剥掉前缀或者加载时用strictFalse并核对缺失项。这个坑不常遇到但一旦遇到很难从 loss 曲线上看出来。5. 进阶用法用 mask 预测做领域适配的验证预训练模型到底有没有学到你领域的东西光看下游 F1 不够直观。我习惯用一个小的验证技巧拿几十条领域句子随机 mask 掉关键实体词让模型做 MLM 预测看 top-5 里有没有合理的候选。这个方法能快速判断继续预训练是否有效比等下游任务跑完再回头排查省时间。import torch def probe_mlm(model, tokenizer, text, topk5): enc tokenizer(text, return_tensorspt) ids enc[input_ids] # 手动把某个位置换成 [MASK]这里演示取第 3 个 token pos 3 masked ids.clone() masked[0, pos] tokenizer.mask_token_id with torch.no_grad(): logits model(masked).logits probs torch.softmax(logits[0, pos], dim-1) top torch.topk(probs, topk) return [tokenizer.decode([i]) for i in top.indices] # 用法对比继续预训练前后的预测合理性 print(probe_mlm(model, tokenizer, 本合同自双方签字之日起生效))这段探测代码的价值在于它把「模型懂不懂领域词」变成一个可观察的输出。如果继续预训练后被 mask 的领域词能稳定进入 top-5说明词表和语义空间确实适配了如果预测全是高频通用词那多半是训练步数不够或者学习率把原知识冲没了。参数上topk取 5 到 10 比较合适太少看不出分布太多噪声大。还有一个值得试的进阶方向是冻结层数实验把编码器底部 6 层冻结只训顶部 6 层加分类头在小样本场景下往往比全量微调更稳过拟合更轻。冻结的实现就是在 optimizer 构造时排除那些层的参数或者训练前把requires_grad置 False。这个技巧在标注数据只有几千条时特别管用我一般会把它和全量微调各跑一遍做对比选验证集更稳的那个。从那以后我每次拿到一份预训练复现包都会先跑一遍 MLM 探测再决定要不要继续预训练而不是上来就堆步数。这个习惯帮我省掉过好几次白跑一整天的训练。希望帮到你。本文还有配套的精品资源点击获取