简介谷歌AI于2018年发布的深度双向Transformer模型BERT在自然语言处理领域树立了预训练语言模型的新标杆。这份实现包采用PyTorch框架完成对该模型的复现代码结构清晰且带有简要注释项目内将模型组件、数据集加载、训练流程与测试入口分模块组织方便读者对照原始论文逐行学习也可直接作为微调和二次开发的基准工程。资源压缩包共包含三十三个文件其中以二十七个Python脚本为核心覆盖模型结构定义、词表与数据迭代器、优化器配置及训练循环等关键环节另有Makefile、配置文件、依赖清单、许可证和说明文档等辅助内容整体体积仅约28KB非常轻量。目前已有878人参与学习与下载适合具有一定PyTorch和NLP基础、希望系统掌握BERT预训练与微调实现细节的研究者和开发者。 2018年BERT刚发布的时候我正被序列标注任务折腾得焦头烂额。那会儿主流预训练思路还是ELMo和GPT-1BERT靠双向Transformer和两个看似简单的预训练任务直接把各种榜单一口气刷穿。问题是官方实现是用TensorFlow写的我那会儿的PyTorch经验更丰富看着一堆estimator和tf.train的代码属实头大。好在没多久GitHub上出现了BERT-pytorch这个项目代码精简到像是把整只BERT拆成了积木我当时花了一个周末跑通又用了两周彻底啃完每一条数据流。这篇文章就是我那份实施笔记的整理会讲BERT的核心原理会拆解PyTorch实现里的关键模块也会手把手带你把环境搭起来、训练一个迷你BERT。无论你是刚接触预训练模型还是想自己动手改模型结构都可以把这篇文章当作一条相对省力的上手路径。1. BERT为什么能在2018年杀出重围双向上下文和两个预训练任务1.1 GPT留下的遗憾单向语言模型的天花板在BERT之前GPT-1已经证明了用Transformer做语言模型预训练再迁移到下游任务是可行的。但GPT用的标准语言模型只能从左往右看每个位置只能用自己的左边信息去预测下一个词。这在生成任务上没毛病可在阅读理解、句子判断这类需要完整理解上下文的场景里就很吃亏。ELMo做了一点改进用两个单向LSTM拼接出伪双向表示但本质还是左右分开没有真正做到同一层里每个位置互相看到彼此。BERT的核心立场很直接与其预测下一个词不如随机盖住一些词让模型同时利用左右两侧的上下文来还原被盖住的内容。这个想法后来也被叫作完形填空它让预训练阶段能真正用上双向信息而不是偷偷把两个方向的向量拼在一起。1.2 Masked Language Model 和 Next Sentence PredictionBERT预训练不像普通语言模型那样只优化一个目标它同时优化两个目标。第一个是Masked Language ModelMLM做法是在输入序列里随机抽取15%的token把这些位置做特殊处理其中80%替换成[MASK]10%替换成一个随机词剩下10%保持原样。模型要还原这些被处理过的token但计算loss时只考虑这15%位置的输出。第二个是Next Sentence PredictionNSP把两句话拼在一起用[CLS]位置上的输出判断第二句话是不是第一句话的下一句。这两个目标一个管词级别理解一个管句子间关系理解合在一起让BERT在各类下游任务上都有了很强的泛化能力。1.3 为什么我们还需要一个PyTorch版本官方代码工程化程度很高但对想搞清楚原理的人来说动辄几千行TensorFlow代码很难下咽。BERT-pytorch这个开源项目反其道而行之把模型部分压缩到几个文件里核心类基本一眼能看完。它没有像HuggingFace那样做大量工业级封装反而更像一份带注释的说明书。我当初就是靠读这个项目才把multi-head attention、embedding、two-task loss这些概念从似懂非懂变成能默写出来。所以这篇博文的另一条主线就是带大家跟着这个项目的设计思路走一遍。2. BERT-pytorch项目代码脉络为什么它看起来这么轻2.1 从数据流水线到训练入口的整体盘一盘BERT-pytorch的代码目录比绝大多数生产级项目都简单大致分这么几块dataset目录负责把原始文本转成BERT需要的输入格式model目录放着embedding、transformer、bert等核心模型类trainer目录负责预训练逻辑。别小看这个轻它的好处是每块代码你都能在半小时内读完整。我当时拿到项目后的读码顺序是先读model/bert.py看整个模型如何调用各组件再读embedding和transformer弄清楚张量在每个层之间的形状变化最后读dataset和trainer搞懂预训练语料怎么变成带[MASK]的样本以及loss怎么算。这样走一遍之后你会对模型前向传播形成一个完整闭环而不是停留在API调用层面。2.2 模型参数配置layer、hidden、head到底意味着什么BERT-base默认配置是12层Transformer Encoder、768维隐藏层、12个注意力头。三层里两头是头数也就是把768维拆成12个64维子空间去做自注意力最后再拼回来。之所以要拆多头是希望不同头能关注不同信息有头抓主语有头抓修饰关系有点像多个小组并行看同一份资料最后一碰头汇总结论。BERT-pytorch在模型初始化时也用类似的参数结构。num_layers、hidden、heads这些值不是摆着好看的它们直接决定了模型参数量。手头计算资源不够的话把12层改成4层、隐藏层改成128维一样能跑通预训练只是效果会差一些。理解这些配置对后面调模型很有帮助。2.3 三段输入token_ids、segment_ids 和 attention_mask用BERT前你必须把原始句子转成三种等长的向量token_ids每个token在词表中的编号segment_ids标记每个token属于第一句话还是第二句话通常第一句是0第二句是1用来支持NSP任务attention_mask标记哪些位置是真实token、哪些是padding出来的。自注意力计算时padding位置会被加一个很大的负偏置让softmax后权重趋近于0等于不让模型关注这些位置。在BERT-pytorch的dataset里还会额外维护一个masked_token_ids把一部分token替换成[MASK]或随机词之后的序列。你看训练代码的时候最容易绕晕的就是token_ids和masked_token_ids的区别前者是干净的原句后者是喂给模型做预测的“残缺句”而loss计算要用前者当标签。3. 亲手拆解核心模块Embedding、多头自注意力、Encoder Layer3.1 三个Embedding加起来的输入矩阵BERT的输入表示不是查一张表就完事而是三张表相加。第一张是token embedding把每个token转成768维向量第二张是segment embedding把句子编号信息也编码进去第三张是position embedding给每个位置一个可学习的向量。理论上可以像Transformer原论文那样用三角函数生成固定的位置编码但BERT用的是可学习的nn.Embedding(max_len, dim)。效果上没有绝对优劣只是可学习版本更灵活。三个向量按位置逐元素相加得到模型的输入。这一步在BERT-pytorch的BERTEmbedding类里实现得很直白先分别跑三个Embedding再执行一次加法。新手容易忽略的是相加之前要确保segment_ids也经过Embedding层别把0/1这种整数直接加进连续向量里。3.2 多头自注意力的计算流程自注意力是整个BERT的核心数学上就一句话Attention(Q, K, V) softmax(Q·Kᵀ/√dₖ)·V。用生活化的比喻来说Q是我在找什么信息K是我身上有什么信息V是我真正想输出的信息。Q和K的点积决定两个位置之间的匹配程度除以√dₖ是为了防止内积数值过大导致softmax梯度饱和再经过softmax归一化权重最后加权求和V。多头版本就是把Q、K、V各自拆成12份分别做12次自注意力最后拼回去线性变换。BERT-pytorch里实现多头注意力的要点是先通过线性层把输入变成Q、K、V然后reshape成(batch_size, num_heads, seq_len, head_dim)交换维度做矩阵乘法加权求和后再reshape回原尺寸。这些reshape动作看起来烦但本质就是张量的分块组合。我读代码时画过一张手动版的流程x经过q_linear、k_linear、v_linear形状是(batch, seq, hidden);把最后一维拆成(head, head_dim)得到(batch, head, seq, head_dim);score q k.transpose(-2, -1) / sqrt(head_dim);对attention_mask中padding位置填充负无穷执行softmax;用attention_weight v得到输出。3.3 从Transformer Layer到整个Encoder一个Transformer Layer由两部分组成多头自注意力 前馈网络。每部分都套着一个残差连接 层归一化的结构。残差连接解决了深层网络退化问题让梯度能更快地回传层归一化则是把每层输出拉回到稳定的均值方差训练起来更顺。BERT-pytorch实现里每个Encoder Layer的伪逻辑如下先做自注意力得到结果后与输入相加并做LayerNorm再过一层带GELU激活函数的两层全连接再残差LayerNorm。这里用的激活函数不是ReLU而是GELUGELU在负半轴保留了少量非零梯度许多实验证明它对Transformer训练有帮助。12层这样的Layer叠起来就是深度双向Transformer Encoder。4. 预训练任务实现中的设计细节为什么这么设计代码怎么写4.1 MLM里80/10/10的比例不是随便定的很多人一开始不理解为什么MLM不能全部用[MASK]。原因是下游微调阶段输入里不会出现[MASK]标记如果预训练时模型过于依赖这个标记迁移时就会“水土不服”。所以BERT设计了一种替代策略80%的概率用[MASK]10%概率替换成随机的其他词10%概率保持原词。这逼着模型不仅要学会看到[MASK]就猜原词还要学会当前词可能是错的要根据上下文判断是否需要修正。在写代码时这个逻辑通常是这样落地的probability torch.rand(seq_len) mask_indices probability 0.15 # 80% - [MASK] # 10% - random token # 10% - original token注意一旦选定了15%的位置原始token一定要存下来当label否则模型学无可学。BERT-pytorch的dataset里就是这么处理的先生成masked_token_ids同时记录mlm_label只在被mask的位置保留原token编号其余位置设为-1计算loss时忽略这些-1位置。4.2 NSP用[CLS]当全句表示Next Sentence Prediction的实现就简单多了。对于每一个训练样本模型接收两个句子A和BB有50%概率是真实的下一句50%概率是从语料里随机抽的。两个句子用[SEP]分隔最前面加[CLS]。经过多层Encoder之后[CLS]位置的输出向量被看作整个序列的聚合表示后面接一个二分类头输出是否是真实下一句的概率。在BERT-pytorch里NSP loss和MLM loss是同时算出来的最后加起来做反向传播。如果你去看它的训练循环会看到它从next_sent这个布尔标签构造二分类loss从masked_label构造词分类loss然后total_loss mlm_loss nsp_loss。4.3 两个loss合在一起会不会互相干扰这是很多人的疑问多任务训练会不会让模型顾此失彼实践证明MLM和NSP这两个任务的目标其实有很强的互补性。MLM迫使模型理解词级别的局部上下文NSP迫使模型学会处理句子间关系。两者共享同一个Transformer主干只在最后各接一个小分类头增加的参数量很小所以主干没有太大负担。BERT-pytorch默认把两个loss简单相加没有做加权也是因为原论文里就是这么做的。实际训练时你可以观察两个loss的量级如果NSP降得特别快而MLM降得慢也不用慌这是正常现象。5. 实操搭建PyTorch环境并训练一个迷你BERT5.1 环境准备conda、Python、PyTorch安装与GPU检查工欲善其事必先利其器。在跑通任何模型之前先得把PyTorch装好。如果你用的是Anaconda推荐用conda创建独立环境避免把基础环境搞乱conda create -n bert_env python3.9 conda activate bert_env pip install torch默认装的PyTorch通常是CPU版还是GPU版取决于你的机器。如果机器有NVIDIA显卡最好按官网指引装对应CUDA版本的PyTorch。装完后一定要做个最基本的检查python -c import torch; print(torch.__version__); print(torch.cuda.is_available())torch.cuda.is_available()返回True才说明GPU可用。如果你只有CPU也能训练迷你BERT就是慢一点建议把本教程的模型调小。5.2 准备一个小型语料并构建预处理流程为了快速看到训练效果我建议先用几段英文新闻或维基百科片段当作语料别一上来就追求大数据量。BERT-pytorch的数据预处理需要把句子切成token并且要构造两句配对和mask逻辑。你可以直接参考项目里的dataset.py或者用下面的思路自己写加载文本文件按换行切出句子用简单的空格分词中文需要先分词这里先以英文为例构造句子对相邻句子作为正例随机句子作为负例转成token编号加上[CLS]和[SEP]并生成segment_ids、attention_mask和masked_input_ids。如果你一时半会不想自己写也可以先直接使用BERT-pytorch仓库里自带的BertDataSet加载一份文本文件它会自动完成这些步骤。当时的项目里有个samples.txt就是干这个用的。5.3 训练参数设置和日志观察迷你BERT我建议这样配hidden64、num_layers2、heads4、seq_len32、batch_size16一句话总结就是能看出效果又不至于等太久。优化器用Adam学习率设成1e-4比原论文的1e-4小一点也没问题。训练若干轮后打印MLM loss和NSP loss。我当时跑出来的曲线大致是前几十个batchMLM loss从10左右快速掉到6NSP loss从0.7掉到0.2再往后下降变缓。看到这种曲线就说明模型在正常学习了。如果你想直接观察效果可以在固定的[MASK]句子上查看模型预测的top词能明显看到模型给出的候选词和上下文是相关的。6. 复现BERT-pytorch时踩过的坑和调优心得6.1 学习率和warmup没有它loss会飘BERT这类Transformer模型对学习率很敏感。我一开始直接上来就是3e-4的固定学习率结果训练不到几百步loss直接变成NaN。后来给优化器加了warmup也就是前10%的step里学习率从0线性增到目标值后面再线性衰减曲线立刻就稳了。BERT-pytorch的NoamOpt类就实现了这种调度用起来非常简单optimizer NoamOpt(hidden, warmup1000, optimizeroptimizer)这个warmup不仅是为了收敛稳定还能避免早期模型还没见到足够样本时就被大梯度带到奇怪区域。6.2 显存爆炸与梯度累积用GPU训练真实BERT-base需要很大的显存至少也得12GB以上试一下。如果你只有8GB甚至更小不要硬刚大batch。实践方案是把batch_size设成2或4然后累计若干步梯度再更新一次参数效果等价于大batch显存压力小得多。核心代码大概是total_loss / accumulation_steps loss.backward() if step % accumulation_steps 0: optimizer.step() optimizer.zero_grad()我第一次跑BERT-pytorch时因为没做梯度累积一张8GB卡直接OOM。把batch_size降到4、再配合梯度累积才顺利跑完一个epoch。6.3 训练完怎么用接一个下游任务预训练本身只是第一步更好玩的是拿这个BERT去改下游任务。你可以把训练好的model提取出来在[CLS]向量后面接一个全连接层做二分类或多分类。不过用迷你BERT微调效果不会太好但它能帮你把整套流程打通。等你将来有条件用HuggingFace加载官方预训练权重再微调时思路是一模一样的。我个人在实际操作中的体会是BERT-pytorch这个项目最宝贵的地方不是它能让你得到一个多强的模型而是它把所有复杂概念都摊开了。读它的过程就像看一个老前辈手写的推导草稿比看那些封装得黑漆漆的工业级代码舒服得多。我后来每次跟新人讲解自注意力都会建议他先把这里的代码自己敲一遍绝对比刷十篇论文有感觉。你如果已经跟着本文走到了这一步下一步不妨试着把隐藏层维度调大一点、再加点真实语料看看loss曲线会不会给你更多惊喜。本文还有配套的精品资源点击获取