BERT-PyTorch源码解析：从注意力机制到预训练全流程

简介：面向NLP学习者与PyTorch使用者，这是Google AI 2018年BERT模型的PyTorch实现，以带注释的简洁代码呈现Transformer双向编码器的预训练思路，可帮助理解语言模型迁移到下游任务的原理。包内共33个文件，27个Python脚本构成核心实现，覆盖模型架构、数据集处理、训练器与入口模块，同时包含Makefile、requirements配置、说明文档等，压缩包仅28KB，轻量且便于逐行研读。发布至今已有878人学习下载，项目代码保持原始结构，配套注释降低了上手门槛，适合想深入BERT原理、动手实践预训练流程的中级NLP开发者，能从中获得从语料准备到模型训练的可运行参考。

1. 项目概述与核心价值

1.1 这个项目到底解决什么问题

先交代一下背景。BERT是Google AI在2018年提出的预训练语言模型，全称是Bidirectional Encoder Representations from Transformers，中文常译作“基于Transformer的双向编码器表示”。当年它一出，直接刷新了11项NLP任务的SOTA成绩，几乎成了自然语言处理领域的“分水岭”式成果。

不过原始官方实现用的是TensorFlow，而PyTorch生态在学术界和工业界的普及度同样极高。于是就有了BERT-pytorch这个项目——把Google官方BERT模型用PyTorch框架重新实现一遍，让PyTorch用户也能跑BERT、微调BERT、甚至从零预训练BERT。

我最初看到这个项目标题时，第一反应是“这不就是换个框架复刻一遍吗？”但真正深入去读代码、跑实验之后才意识到，事情没这么简单。BERT-pytorch不是一个简单的“翻译”项目，它把BERT的完整结构拆解成了清晰可读的PyTorch模块，包括Tokenizer、Embedding、Transformer Encoder层、预训练任务（NSP + MLM）等，每一步都给了简洁的实现。对想搞懂BERT内部原理的人来说，这个项目的代码比官方TensorFlow实现好读太多了。

1.2 适合谁读、能带来什么收获

这个项目的目标读者非常明确：一类是刚入门NLP、想搞懂Transformer和BERT内部机制的开发者，另一类是手里只有PyTorch环境、想把BERT用起来的研究生或工程师。前者可以从代码里学到BERT的核心组件怎么拼装，后者可以直接基于这份实现做下游任务微调。

我个人的体感是，如果你已经把transformers库用得滚瓜烂熟，再来看这个项目会有一种“原来封装底下的东西长这样”的通透感。因为transformers库为了兼容所有BERT变体，把代码抽象得比较深，有时候反而不容易看清主脉络。BERT-pytorch则把注意力机制、多头自注意力、LayerNorm、位置编码、预训练损失函数全部摊开在面前，每一行都能对着论文找到来源，非常适合作为“BERT原理阅读的配套代码”。

这篇文章我就基于自己踩坑、调试、二次开发的经验，把这个项目从环境搭建、代码结构、关键实现细节到常见坑位全部捋一遍，希望能帮你少走弯路。

2. 环境准备与PyTorch版本选型

2.1 推荐的环境组合与安装步骤

先把环境搞定。BERT-pytorch的代码不算新，对PyTorch版本没有特别苛刻的要求，但我实测下来，不同版本组合的表现差异还是有的。先说结论，我最后稳定跑通的组合是：

• Python 3.9或3.10
• PyTorch 2.1.x（CPU或CUDA版均可）
• torchtext（注意版本，这个项目里做WordPiece Tokenizer时依赖了torchtext的BasicEnglish等工具，但新版torchtext API变化很大）
• CUDA 11.8或12.1（如果要用GPU）

安装PyTorch时，我建议直接用官方命令生成器。比如要装CUDA 12.1版本：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

如果是CPU版本，直接：

pip install torch torchvision torchaudio

这里有个容易踩的坑：如果你用Anaconda管理环境，千万别在base环境里硬装，一定要先建一个干净的环境：

conda create -n bert_env python=3.9 -y conda activate bert_env

我一开始图省事，直接往base环境里塞了一堆包，结果torchtext和torch版本冲突，报了一堆奇奇怪怪的错。新建环境之后，世界清净了。

2.2 关于CUDA、显卡驱动和torchtext的版本适配问题

GPU版本的PyTorch安装，最容易出问题的就是“驱动版本和CUDA版本不匹配”以及“PyTorch要求的CUDA版本和本机CUDA版本不一致”。

这里分享一个核心认知：PyTorch的CUDA版本并不需要和你系统里安装的CUDA Toolkit完全一致，PyTorch是自带CUDA runtime的。你只需要保证显卡驱动够新就行。比如你想装cu121的PyTorch，但系统里CUDA Toolkit还是11.8，这完全没问题，只要驱动版本支持CUDA 12.1就行。怎么查驱动支持的最高CUDA版本？在命令行敲nvidia-smi，右上角会显示“CUDA Version: 12.1”之类的，这代表你的驱动最高支持到12.1，只要PyTorch的CUDA版本号小于等于这个数就行。

torchtext是个比较麻烦的东西。BERT-pytorch项目里的tokenizer.py使用了torchtext.datasets.text_classification或者torchtext.data.utils中的tokenizer。老版本torchtext（0.9及以下）有torchtext.data.functional.sentencepiece_numericalizer等接口，新版本（0.12+）API大改，很多老代码直接跑不动。如果你想用BPE或WordPiece tokenizer，更推荐直接用tokenizers库或者transformers库自带的Tokenizer，绕开torchtext的版本坑。

我自己在复现时，直接改用了transformers库的BertTokenizer来做tokenization，因为BERT-pytorch官方训练用的vocab文件本质上就是WordPiece格式，和BertTokenizer完美兼容。这样既享受了transformers的便利，又保持了项目的核心训练逻辑不变。

3. 核心代码结构拆解与实现原理

3.1 从目录结构看项目脉络

把项目clone下来之后，先看目录结构：

BERT-pytorch/ ├── bert_pytorch/ │ ├── __init__.py │ ├── __main__.py │ ├── dataset/ │ │ ├── __init__.py │ │ ├── dataset.py │ │ ├── tokenization.py │ │ └── vocab.py │ ├── model/ │ │ ├── __init__.py │ │ ├── attention.py │ │ ├── bert.py │ │ ├── embedding.py │ │ ├── layer_norm.py │ │ ├── linear.py │ │ ├── transformer.py │ │ └── utils.py │ └── trainer/ │ ├── __init__.py │ ├── optim_schedule.py │ ├── pretrain.py │ └── trainer.py ├── scripts/ │ ├── preprocess.py │ └── train.py └── tests/ └── sanity_check.py

这个结构非常清晰，model目录下的每个文件几乎对应Transformer论文里的一个核心概念。transformer.py里定义了TransformerBlock，里面包含多头自注意力（MultiHeadedAttention）和位置前馈网络（PositionwiseFeedForward）。attention.py则是自注意力的具体实现。embedding.py负责Token Embedding、Segment Embedding、Positional Embedding的构造。bert.py则是最终模型的组装。

3.2 关键组件：Embedding层的实现细节

BERT和原始Transformer的一个重要区别在于输入表示。BERT的输入由三部分加和而成：

• Token Embedding：把每个词映射成768维向量（BERT-base配置）
• Segment Embedding：区分两个句子，句子A和句子B的segment id分别为0和1
• Positional Embedding：这里用的是可学习的绝对位置编码，而不是Transformer原论文里的正弦余弦函数

这段逻辑在embedding.py里实现得很直白：

class BERTEmbedding(nn.Module): def __init__(self, vocab_size, d_model, n_segments, max_len, dropout=0.1): super().__init__() self.token_embedding = nn.Embedding(vocab_size, d_model) self.segment_embedding = nn.Embedding(n_segments, d_model) self.position_embedding = nn.Embedding(max_len, d_model) self.dropout = nn.Dropout(p=dropout) def forward(self, tokens, segments): seq_len = tokens.size(1) positions = torch.arange(seq_len, device=tokens.device).unsqueeze(0).expand_as(tokens) embeddings = self.token_embedding(tokens) \ + self.position_embedding(positions) \ + self.segment_embedding(segments) return self.dropout(embeddings)

这里有一个值得品味的点：position_embedding用的是nn.Embedding，意味着位置编码是“学出来的”，而不是固定的。在原始Transformer论文里，作者用固定的三角函数位置编码，后来发现可学习的位置编码在类似规模的数据下效果也够好，BERT就选择了可学习方案。如果你自己实现时候想把sinusoidal位置编码换进去，直接在forward里生成常数矩阵就行，但要注意和后续预训练阶段的分布保持一致。

3.3 自注意力与多头机制的完整解读

attention.py实现了缩放点积注意力（Scaled Dot-Product Attention）。核心公式是：

[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V ]

代码里也就是这几行的事：

def forward(self, query, key, value, mask=None): scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) p_attn = F.softmax(scores, dim=-1) return torch.matmul(p_attn, value), p_attn

这里面关键的“为什么除以根号d_k”值得多说一句。如果不缩放，当d_k比较大时，点积结果的方差会变大，导致softmax之后的梯度非常小，模型就很难训练。把点积除以(\sqrt{d_k})，让方差回到1量级，softmax的输入分布才比较稳定。这是原论文作者通过实验验证出来的细节，很多初学者容易忽略，但实际上对训练稳定性影响很大。

多头注意力（Multi-Head Attention）的实现也很干净：把Q、K、V通过线性变换投影到多个子空间，并行计算注意力头，最后拼接并再投影一次。多头的好处是让模型能够同时关注不同位置、不同表征子空间的信息，比如一个头可能关注语法关系，另一个头关注指代信息。代码里通过view和transpose完成头的拆分和合并，不算复杂，但动起手自己写一遍还是很考验基本功的。

3.4 Transformer Encoder层的组装逻辑

transformer.py定义了TransformerBlock，里面是标准的两层结构：

1. 多头自注意力子层，后接残差连接和LayerNorm
2. 位置前馈网络子层（两层线性+GeLU激活），后接残差连接和LayerNorm

class TransformerBlock(nn.Module): def __init__(self, hidden, attn_heads, feed_forward_hidden, dropout): super().__init__() self.attention = MultiHeadedAttention(h=attn_heads, d_model=hidden) self.feed_forward = PositionwiseFeedForward( d_model=hidden, d_ff=feed_forward_hidden, dropout=dropout) self.input_sublayer = SublayerConnection(size=hidden, dropout=dropout) self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout) def forward(self, x, mask): x = self.input_sublayer(x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask)) x = self.output_sublayer(x, self.feed_forward) return x

这里使用了SublayerConnection来封装残差连接和LayerNorm，顺序是x + sublayer(LayerNorm(x))，也就是Pre-LN的模式。这个细节值得注意：Transformer原论文用的是Post-LN（先加残差再做LayerNorm），而BERT用的是Pre-LN。Pre-LN在训练早期更稳定，对学习率没那么敏感，这也是BERT能够用较大学习率预训练的原因之一。

BERT-base的整体配置是12层TransformerBlock，hidden size=768，12个注意力头，feed-forward中间层维度是3072。你可以通过修改这些超参来拼出不同规模的模型。注意，hidden size必须能被注意力头数整除，否则view操作会报错。这是非常常见的一个低级错误。

4. 预训练任务与训练流程全解析

4.1 MLM与NSP两个预训练目标的设计逻辑

BERT的预训练阶段有两个任务，缺一不可。

第一个是Masked Language Model（MLM，掩码语言模型）。做法是随机遮盖句子中15%的token，然后让模型去预测被遮盖的token。但这里有一个小细节：在这15%的token里，不是全部用[MASK]替换，而是有讲究的——其中80%换成[MASK]，10%换成随机词，剩下10%保持不变。为什么要这样？因为如果训练时看到的全是[MASK]，而下游任务中根本不会出现[MASK]，就会产生预训练和微调之间的分布不一致。加入一定比例的随机词和原词，能迫使模型更多依赖上下文来预测，而不是单纯记住[MASK]标记本身。这个小trick在代码里体现为一个随机选择逻辑。

第二个是Next Sentence Prediction（NSP，下句预测）。输入是一对句子A和B，50%的情况下B是A的下一句（标签为1），50%的情况下B是随机句子（标签为0）。模型需要判断两句是否连续。这个任务让BERT学会句子级别的关联信息，对问答、自然语言推理等任务帮助很大。后来的实验中，有论文对NSP是否必需提出过质疑，比如RoBERTa模型就移除了NSP，但在BERT的本体设计里，NSP和MLM共同组成了完整的预训练目标。

4.2 数据准备：从文本到训练样本的完整流水线

在scripts/preprocess.py里，原始文本要经过一系列处理才能变成BERT的输入格式：

1. 文本清洗，去掉特殊符号，统一小写（是否区分大小写可配置）
2. WordPiece分割，把单词切成子词，比如“playing”切成“play”和“##ing”
3. 加上[CLS]和[SEP]标记，[CLS]放在序列开头，[SEP]放在每个句子结尾
4. 构建token ids、segment ids、attention mask（如果需要的话）
5. 按照最大长度做padding或截断

这里的核心工具是Vocabulary，它维护一个“词和id”的映射表。BERT-pytorch默认提供了一个简单的vocab构建方式，但如果你要复现BERT-base的效果，直接使用Google发布的全词表bert-base-uncased-vocab.txt更省事，有30522个词条。

在实际操作中，我建议先写一个小脚本验证一下tokenize的结果：

from transformers import BertTokenizer tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") tokens = tokenizer.tokenize("Hello, how are you playing today?") print(tokens) # ['hello', ',', 'how', 'are', 'you', 'playing', 'today', '?']

这样基本就能和BERT-pytorch的输入对齐。

4.3 Trainer与前向传播的最小实现

项目的trainer/pretrain.py是预训练的核心循环。整个流程并不复杂：

for epoch in range(epochs): for batch in data_loader: token_ids, segment_ids, masked_tokens, masked_pos, is_next = batch logits_lm, logits_cls = model(token_ids, segment_ids, masked_pos) loss_lm = criterion(logits_lm.transpose(1, 2), masked_tokens) loss_cls = criterion(logits_cls, is_next) loss = loss_lm + loss_cls optimizer.zero_grad() loss.backward() optimizer.step()

注意这里的logits_lm并不是对所有token都计算损失，而是只对被mask的位置计算。masked_pos记录了每个样本中被mask的位置索引，模型根据这些索引去取对应的输出向量，再和真实token_ids做交叉熵。这个细节过滤掉了很多无关位置的噪声，是模型能够高效学习的关键。

学习率调度上，BERT-pytorch实现了Transformer论文里的warmup + decay策略。前若干步学习率线性上涨，之后逐步衰减，这样做能避免早期更新幅度太大导致训练不稳定。optim_schedule.py里自己实现了这个调度器。如果你用PyTorch 2.x，也可以直接用现成的get_cosine_schedule_with_warmup等接口替代，效果类似。

4.4 从零开始训练的最小实验记录

为了验证整个流程，我用一个小型数据集（约50MB英文新闻文本）做了从零开始的预训练实验，配置为4层Transformer、hidden size=256、4个注意力头。这个配置远小于BERT-base，但足以跑通流程。训练的batch size设为32，序列长度128，Adam学习率1e-4，warmup步数1000。

跑了大约3万步之后，MLM的准确率从初始的不到1%涨到了35%左右，NSP准确率稳定在80%以上。如果用标准的BERT-base配置去跑完整预训练，需要16张TPU跑4天，所以自己玩的话一定要调小参数。这个“调参”的过程本身就是理解BERT的好机会——你随便改一个参数，都能直观感受到训练速度和loss的变化，这是直接用官方预训练模型体会不到的。

5. 常见问题与避坑手记

5.1 安装和依赖问题速查

这是新手最容易卡住的环节，我整理了几个高频问题：

问题1：torchtext导入报错

老版本BERT-pytorch的tokenization.py依赖torchtext的bert vocab相关接口。新版torchtext里很多API被移除，直接导入会报AttributeError。解决办法是改用transformers库的tokenizer，或者把torchtext锁在0.9.0版本使用。我更推荐前者，因为transformers生态更持续。

问题2：训练时显存不足（OOM）

BERT-base输入长度512、batch size 16在GTX 1080Ti上就非常吃力。解决办法：

• 降低序列长度到128或256
• 减小batch size
• 打开混合精度训练torch.cuda.amp.autocast()

问题3：masked_fill中的-1e9会导致loss为NaN吗？

如果mask把整行都遮住了，softmax后全0，反向传播时会出现NaN。这种情况在padding部分可能出现。更安全的做法是用torch.finfo(dtype).min代替-1e9，并确保每一个样本中至少有一个位置是有效的。

5.2 训练过程中的几个反直觉经验

有些经验光看文档是学不到的，必须实际操作过才能体会：

第一，输入序列的长度对训练速度影响是线性的，但对内存影响是平方级的。因为注意力分数的形状是[batch, heads, seq_len, seq_len]，序列长度从128提升到256，内存消耗直接变成原来的4倍（256/128的平方）。所以做实验时先从短序列跑通，再逐步加长。

第二，LayerNorm层的epsilon参数不要随便动。项目中用的是1e-12，这个值看起来小得离谱，但BERT官方实现里就是这样的。如果你把它改成常见的1e-5，反而可能导致训练初期不稳定。这个细节也是从官方代码里传下来的“奇奇怪怪的设定”。

第三，预训练的warmup步数不能太少。我在小数据集上踩过坑，warmup只设100步，结果前2000步loss大幅震荡，后来把warmup加到1000步就稳定了。要知道，在大型语料上，warmup步数经常要设置到10000步以上。

第四，如果条件允许，直接用Google发布的预训练权重做微调，而不是从零预训练。从零预训练需要的数据量和算力对个人开发者来说非常不友好。BERT-pytorch的意义更多在于“理解原理”和“复现实验”，而不是和生产环境里用现成权重抢精度。做下游情感分类、命名实体识别等任务时，加载bert-base-uncased权重再微调，效果和速度会好得多。

5.3 一个实用的复现检查清单

如果你想完整跑通BERT-pytorch并验证自己的理解，我建议按这个顺序检查：

1. 输入tokens是否正确添加了[CLS]和[SEP]，长度是否一致
2. 多头注意力的num_heads是否能够整除d_model
3. Embedding层三个矩阵的维度是否都是vocab_size, d_model/max_len, d_model/n_segments, d_model
4. MLM损失是否只统计masked位置，而不是所有位置
5. NSP的标签是否和segment id满足对应关系（“句子B确实是随机句或下一句”）
6. 训练时有无用一个小batch做一次前向和反向，确保没有维度不匹配的报错

这个清单看起来非常基础，但我敢说至少有50%的人第一次复现时报错都源自这几个地方。我自己当时就在多头注意力拆分合并那一步卡了很久，最后打印每个tensor的shape才排查出来。

6. 基于BERT-pytorch的二次开发思路

6.1 如何替换成你自己的Tokenizer

BERT-pytorch默认的tokenization模块比较简陋，如果你要做中文文本的处理，建议直接替换成transformers的BertTokenizer。中文BERT有两种做法：一种是基于字（character）切分，每个汉字作为一个token；另一种是基于词切分。Google提供的中文模型bert-base-chinese就是按字切分的。替换时只需要把原始的tokenizer换成：

from transformers import BertTokenizer tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")

然后改写数据预处理流程里调用tokenizer的部分，保持输出的token_ids、segment_ids仍是list of int即可。后续模型代码完全不用动。

6.2 如何用这个项目做下游任务微调

这里提供一个典型的文本分类微调思路。先加载BERT-pytorch的预训练模型（可以是Google权重转换过来的），然后取出[CLS]位置的输出向量，接一个全连接分类层：

class BertClassifier(nn.Module): def __init__(self, bert_model, num_classes): super().__init__() self.bert = bert_model self.classifier = nn.Linear(bert_model.hidden, num_classes) def forward(self, tokens, segments): _, pooled = self.bert.forward(tokens, segments) # 这里需要从输出中取出CLS位置的hidden state logits = self.classifier(pooled) return logits

注意，BERT-pytorch模型的bert前向返回的是所有token的隐藏状态（[batch, seq_len, hidden]），你要取[:, 0, :]才是[CLS]的表示。很多代码里所谓的“pooled output”其实就是一个简单的[CLS]取出操作，不像transformers库中还有额外的pooler层。这个细节在微调时直接影响分类头的效果。

微调的技巧也很直接：用较低学习率（2e-5到5e-5），因为预训练好的参数已经处于一个较优区域，学习率太大容易破坏它。通常只需要跑3~5个epoch就能在下游任务上得到不错的效果。

6.3 项目之外：从BERT到更现代的模型

BERT-pytorch虽然经典，但如今NLP领域已经有了很多更强的模型。但坦白说，如果你想真正理解RoBERTa、ALBERT、DistilBERT，乃至GPT系列的核心思想，先啃透BERT-pytorch的代码会是一个非常扎实的起点。

因为后续很多模型都是在BERT的骨架上做减法或加法——RoBERTa去掉了NSP但加大了数据和训练时长，ALBERT用参数共享和矩阵分解大幅度减少了参数量，DistilBERT用知识蒸馏把模型缩小了40%。你理解了BERT的每一块部件，再去看这些模型的“改动点”就会非常轻松。

我在实际做项目时，也经常会把BERT-pytorch的代码拿来当作“最小可运行”的基线，用来测试新数据集上是否存在数据泄漏、处理流程是否合理。它的代码足够短，短到你可以在半小时内读懂全部关键逻辑，这一点在生产级的transformers库中是做不到的。

7. 最后的实操心得

这篇文章写到这里，其实已经没有“总括性结论”的必要了。我记得自己第一次独立把BERT-pytorch从环境搭建到预训练小模型跑通，花了整整一个周末。期间踩过torchtext的坑，踩过CUDA版本不对的坑，也踩过多头注意力维度错误导致loss不下降的坑。但正是这些排查过程，让我对BERT的理解远超“会用transformers库”的程度。

如果让我给刚接触这个项目的朋友三条建议，第一，先别急着跑大数据，拿一个只有几千条的小语料，把整个训练流程跑通，观察loss变化，理解每一步在干什么；第二，一定要动手改参数，改hidden size、改层数、改学习率，看看训练曲线有什么不同，这种直观感知是看多少文章都替代不了的；第三，把官方预训练权重转换到PyTorch环境后用在下游任务上，这样既能保证效果，又能体会BERT-pytorch代码和官方权重的兼容性。

这个项目虽然名叫“BERT-pytorch”，但它实际上是一个深度学习爱好者极佳的学习素材——它把一个曾经刷新无数榜单的模型，压缩到了几千行清晰可读的PyTorch代码里。弄懂它之后，你再看其他NLP模型的实现，大概率都能举一反三。这大概就是经典项目给人留下的最宝贵的财富。

本文还有配套的精品资源，点击获取