★ FEATURED ARTICLE 简介面向NLP开发者与PyTorch学习者的BERT实现资源对应Google AI 2018年发布的BERT预训练语言模型基于深度双向Transformer架构适用于文本分类、问答、序列标注等自然语言处理任务。代码采用简洁注释组织模块划分清晰便于对照原论文逐行研读尤其适合刚接触Transformer架构的中级开发者。压缩包共33个文件以27个Python脚本为主辅以Makefile、requirements.txt、YAML配置与README等工程化文件整体仅28KB轻量易读。已有879人浏览学习。资源内含模型定义、数据集处理、训练器与主入口等核心模块并配有安装命令与两行制语料准备示例可帮助读者快速跑通预训练流程理解自注意力机制、掩码语言模型与下一句预测等关键训练目标。通过阅读源码还能掌握分词处理、批次构造与优化器设置等实践细节便于后续微调下游任务。需注意项目仍在完善中部分代码未经严格验证用于学习研究较为合适。1. BERT-pytorch 项目解析一个可直接复现的 Transformer 预训练模型在 NLP 领域2018 年发布的 BERT 是一个绕不开的里程碑。这个 BERT-pytorch 项目把原版 TensorFlow 实现迁移到了 PyTorch 框架且保留了关键的 BERT-base 结构12 层 Transformer encoder、768 维隐藏层、12 个 attention head总参数量约 1.1 亿。它不是魔改版也不建议拿它去跑极大规模预训练它的价值在于把 BERT 的核心结构剥开让开发者用 PyTorch 原生 API 直观地看到语言模型language model是怎么搭建的并可快速二次开发。如果你正在学习 transformer、pytorch 的底层机制或者想基于预训练模型做下游任务微调这份源码比任何教程都来得直接——因为它本身就是可编译、可训练的最小实现。2. 技术栈与整体架构PyTorch 下的 Transformer 完全体2.1 技术选型为什么是 PyTorch 而不是原版 TF原版 BERT 是基于 TensorFlow 1.x 实现的代码里有大量tf.train、tf.contrib等旧 API 依赖想在自己的环境中跑起来非常麻烦。这个项目选择 PyTorch核心原因有三点一是 PyTorch 的动态图机制让模型代码更接近公式本身nn.Module的写法对阅读者极其友好二是 HuggingFace 生态流行后PyTorch 已经成为 NLP 开发的事实标准三是这个项目刻意不额外引入transformers这种高级库所有模块都基于torch.nn手工实现方便二次开发时按 needs 增删组件。技术栈可以看出项目非常「轻依赖」除了torch和numpy几乎没有必须安装的第三方包。数据处理时用到tokenizers或者tqdm也属于 optional。这种极简依赖让项目在任何有 GPU 或纯 CPU 的机器上都能复现。模块依赖说明模型层torch.nn, torch.nn.functional手工实现所有 transformer 组件数据处理torch.utils.dataDataset、DataLoader、RandomSampler优化器torch.optimAdamW、SGD支持 warmup额外工具numpy, tqdm数据矩阵转换、进度显示2.2 模块划分一份清晰到可以直接抄的模型结构项目代码结构紧凑核心都集中在bert_pytorch/包下目录并不复杂bert_pytorch/ ├── __init__.py ├── model/ │ ├── __init__.py │ ├── embedding/ │ │ ├── token_embedding.py │ │ ├── segment_embedding.py │ │ ├── position_embedding.py │ │ └── bert_embedding.py │ ├── attention/ │ │ └── attention.py │ ├── encoder/ │ │ ├── encoder.py │ │ ├── encoder_block.py │ │ ├── feed_forward.py │ │ └── transformer.py │ └── utils/ │ ├── feed_forward.py │ ├── layer_norm.py │ ├── gelu.py │ └── mask.py ├── dataset/ │ ├── __init__.py │ ├── dataset.py │ ├── tokenizer.py │ └── vocab.py ├── trainer/ │ ├── __init__.py │ └── pretrain.py └── __main__.py每个子模块都对应 BERT 论文中的一个具体组件。其中embedding目录下包含了三种 embedding 的独立实现token embedding 是词表映射矩阵segment embedding 区分句子 A / Bposition embedding 使用正弦余弦函数生成位置编码。三个独立类最后在bert_embedding.py里汇总对应论文中的E E_token E_segment E_position。attention目录中实现的是多头的 scaled dot-product attention属于整个模型最核心的部分。encoder目录则把 12 层 encoder block 堆叠起来每一层包含 self-attention、feed forward、layer normalization 和残差连接。整个模型结构就是把这些类按论文原顺序拼装起来。3. 核心代码逐模块拆解从 Embedding 到 Pretrain 的完整链路3.1 Token Embedding 与 Segment Embedding词表映射和句子标记BERT 输入有两类特殊标记[CLS]放在句首[SEP]分隔两个句子。token embedding 做的就是最常规的词表映射把 tokenizer 切出来的 token id 映射到 d_model 维度的向量。class TokenEmbedding(nn.Embedding): def __init__(self, vocab_size, d_model): super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx0)这个类直接继承nn.Embedding内部就是一个查表操作。padding_idx0意味着词表中索引为 0 的 token通常是[PAD]在梯度更新时会被忽略避免 padding 部分干扰训练。vocab_size在你的语料上训练 tokenizer 后确定常见值像 30522BERT-base 原始词表。Segment embedding 稍微不一样它要区分两个句子。所以SegmentEmbedding是一个 embedding 矩阵但 vocab_size 固定为 2——type 0 代表句子 Atype 1 代表句子 Bclass SegmentEmbedding(nn.Embedding): def __init__(self, embed_size): super(SegmentEmbedding, self).__init__(2, embed_size, padding_idx0)这里embed_size就是d_model768。输入给它的 tensor 是[batch, seq_len]的 0 / 1 序列输出[batch, seq_len, 768]的段向量。注意padding_idx0在这个类里表示 segment 标记为 0 的位置也是句子 A 的一部分不是真的 padding这不影响使用。3.2 Positional Embedding正弦函数位置编码的实现BERT 没有 RNN / CNN 这种自带顺序的结构所以必须显式注入位置信息。项目里的位置编码沿用了 transformer 论文里的正弦余弦函数方案class PositionalEmbedding(nn.Module): def __init__(self, d_model, max_len512): super(PositionalEmbedding, self).__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) pe pe.unsqueeze(0).transpose(0, 1) self.register_buffer(pe, pe) def forward(self, x): return x self.pe[:x.size(0), :]div_term的计算是这段代码的数学关键对于输入的第i个维度如果 i 为偶数频率因子是10000^(-i/d_model)奇数对应余弦。register_buffer把pe注册为模型持久化 buffer不会参与梯度更新但会随模型保存到 state_dict 里。在 forward 里x.size(0)是序列长度直接切出前seq_len行加到 embedding 上。3.3 Multi-Head Self-Attention注意力机制的核心实现这一块是 BERT 的心脏。项目把 attention 两层封装MultiHeadAttention外部做线性变换和头拆分内部一个ScaledDotProductAttention做真正 attention 计算。class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_head, d_k, d_v): super(MultiHeadAttention, self).__init__() self.n_head n_head self.d_k d_k self.d_v d_v self.w_qs nn.Linear(d_model, n_head * d_k) self.w_ks nn.Linear(d_model, n_head * d_k) self.w_vs nn.Linear(d_model, n_head * d_v) self.fc nn.Linear(n_head * d_v, d_model) def forward(self, q, k, v, maskNone): d_k, d_v, n_head self.d_k, self.d_v, self.n_head sz_b, len_q, _ q.size() sz_b, len_k, _ k.size() sz_b, len_v, _ v.size() q self.w_qs(q).view(sz_b, len_q, n_head, d_k).transpose(1, 2) k self.w_ks(k).view(sz_b, len_k, n_head, d_k).transpose(1, 2) v self.w_vs(v).view(sz_b, len_v, n_head, d_v).transpose(1, 2) attn, output self.scaled_dot_product_attention(q, k, v, mask) output output.transpose(1, 2).contiguous().view(sz_b, len_q, -1) output self.fc(output) return output, attnforward 中先通过三个线性层得到 Q、K、V然后view成[batch, len, n_head, d_k]transpose(1, 2)切出头维度变成[batch, n_head, len, d_k]。attn是每头计算出的注意力权重矩阵output是加权聚合后的向量最后fc把所有头拼接后的向量投影回d_model维度。这里contiguous()很重要transpose后的 tensor 内存布局不连续直接view会报错。真正的注意力计算封装在ScaledDotProductAttentionclass ScaledDotProductAttention(nn.Module): def __init__(self, d_k, attn_dropout0.1): super(ScaledDotProductAttention, self).__init__() self.dropout nn.Dropout(attn_dropout) def forward(self, q, k, v, maskNone): scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn F.softmax(scores, dim-1) attn self.dropout(attn) output torch.matmul(attn, v) return output, attnmasked_fill把 mask 中 0 的位置填成 -1e9这样在 softmax 之后这些位置的注意力权重会趋近于 0实现对 padding 或未来信息的屏蔽。在 BERT 预训练里mask 是 1 或 0 的张量1 表示需要参与计算的位置。3.4 Position-wise Feed-Forward 与 LayerNorm残差块的另一半每层 transformer block 中attention 之后是一个两层全连接网络class PositionwiseFeedForward(nn.Module): def __init__(self, d_in, d_hid, dropout0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 nn.Linear(d_in, d_hid) self.w_2 nn.Linear(d_hid, d_in) self.layer_norm nn.LayerNorm(d_in, eps1e-6) self.dropout nn.Dropout(dropout) def forward(self, x): residual x x self.w_2(torch.relu(self.w_1(x))) x self.dropout(x) x x residual x self.layer_norm(x) return xd_in是 768d_hid是 3072这是论文里d_ff 4 * d_model的体现。激活函数用的是 Relu实际 BERT 论文里是 GELU但项目在gelu.py里单独实现了需要用到的 GELU这个 Feed-Forward 里默认用 relu具体阅读时注意这个差异。残差连接先加后 norm和EncoderBlock里先 norm 后 attention 的顺序不同两处 LayerNorm 的位置并不完全一致运行上不影响整体。3.5 EncoderBlock 与堆积逻辑12 层如何组合每个 encoder block 是一个完整的 transformer 单元class EncoderBlock(nn.Module): def __init__(self, d_model, n_head, d_k, d_v, d_ff, dropout0.1): super(EncoderBlock, self).__init__() self.slf_attn MultiHeadAttention(d_model, n_head, d_k, d_v) self.pos_ffn PositionwiseFeedForward(d_model, d_ff, dropout) self.dropout nn.Dropout(dropout) def forward(self, enc_input, slf_attn_maskNone): residual enc_input enc_input, attn self.slf_attn(enc_input, enc_input, enc_input, maskslf_attn_mask) enc_input self.dropout(enc_input) enc_input enc_input residual enc_input nn.functional.layer_norm(enc_input, enc_input.size()[1:]) enc_output self.pos_ffn(enc_input) return enc_output, attn这里的残差连接并没有调用PositionwiseFeedForward内部的 layer_norm而是直接调nn.functional.layer_norm输入维度是整个特征维度enc_input.size()[1:]。Transformer模块则负责循环堆叠 12 个这样的 block并且保证输出连续 tensor 以便后续任务使用。3.6 预训练任务Masked LM 与 Next Sentence Prediction预训练实现集中在trainer/pretrain.py包括两个任务。Masked LM 的构造方式是训练前按 15% 概率随机把 token 替换为[MASK]模型输出经过每个 token 位置的 768 维向量过一个线性层映射回词表大小再用交叉熵计算 lossmasked_loss self.criterion(masked_output, masked_label) # criterion torch.nn.CrossEntropyLoss(ignore_index0)ignore_index0表示 padding 位置不参与 loss 计算。Next Sentence Prediction 则是一个二分类任务判断输入的第二句是否真实衔接第一句复用[CLS]位置的输出向量过一个二分类线性层计算 loss。两个 loss 相加后进行反向传播这正是 BERT 预训练阶段赖以学习句法、语义和上下文依赖的关键途径。4. 数据准备与训练脚本把项目在本地跑起来的完整流程4.1 数据格式要求与 Dataset 实现项目的训练数据默认支持两种格式纯文本文件每行一句话和符合特定要求的文件格式每一行包含[CLS] token... [SEP] token... [SEP]的原始文本。数据集类在dataset/dataset.py里做两处核心处理一是把连续句子以(sentence1, sentence2, is_next)三元组形式打包二是构造 masked LM 的输入 id 和 label。def __getitem__(self, idx): t1, t2, is_next self.data[idx] tokens self.tokenizer.tokenize(t1) [[SEP]] self.tokenizer.tokenize(t2) [[SEP]] ... # 随机 mask 15% 的 token input_ids [self.vocab.stoi[token] for token in tokens] segment_ids [0] * (len(self.tokenizer.tokenize(t1)) 2) [1] * (len(self.tokenizer.tokenize(t2)) 1) return torch.tensor(input_ids), torch.tensor(segment_ids), torch.tensor(is_next), torch.tensor(masked_label)segment_ids的构造很关键句子 A 的所有 token 加上[SEP]标记为 0句子 B 加上[SEP]标记为 1。__getitem__返回四个 tensor输入 id、segment id、next sentence 标签、masked token 标签。这些 tensor 最终通过 DataLoader 组装成 batch。4.2 Tokenizer 与 Vocab 构建中文英文都能用项目的tokenizer.py没有用 BPE 之类的复杂算法而是用字典序列的方式class Tokenizer: def __init__(self, vocab): self.vocab vocab def tokenize(self, text): import re words re.findall(r\w|[^\w\s], text) return words def convert_tokens_to_ids(self, tokens): return [self.vocab[token] for token in tokens] def convert_ids_to_tokens(self, ids): return [self.ids_to_tokens[id] for id in ids]正则表达式\w可以匹配连续字母数字或中文[^\w\s]匹配标点符号。这也意味着你可以直接喂中文文本\w在 unicode 模式下可以匹配中文字符tokenizer 会把每个连续中文片段视为一个 token。4.3 训练命令与关键参数调整要跑预训练命令行入口在bert_pytorch/__main__.py典型命令如下python -m bert_pytorch --train_data /path/to/corpus.txt \ --vocab /path/to/vocab.json \ --output_path /path/to/model_model.bin \ --hidden 768 \ --n_layers 12 \ --attn_heads 12 \ --seq_len 64 \ --batch_size 16 \ --epochs 5 \ --lr 1e-4 \ --mask_prob 0.15常用参数含义参数默认值 / 建议值说明--hidden768隐藏层维度必填--n_layers12encoder 层数--attn_heads12attention head 数量--seq_len64最大序列长度超过会截断--batch_size16训练 batch 大小显存不够就调小--mask_prob0.15Masked LM 的 mask 概率论文标准值如果你是新手第一次跑建议把seq_len设成 64、batch_size设成 8用纯 CPU 也能几分钟内跑通一次迭代。若想直接体验前向传播可以用--eval或写一个简单脚本只加载模型输出 logits。4.4 加载预训练权重或自己继续训练的两种路径项目支持的模型保存和加载方式有两种。一种是完全重新训练即上面命令生成的 model.bin加载时model.load_state_dict(torch.load(path))即可。另一种是如果想要基于别人训练好的 BERT 权重做微调常见做法是加载 HuggingFace 的BertModel权重然后转置键名映射到本项目结构里由于本项目参数命名与 HuggingFace 不完全一致需要自己写一个映射脚本这一步对二次开发来说属于家常便饭。5. 避坑指南我在复现过程中踩过的五个关键坑5.1 训练时 loss 变成 NaN现象训练到几百步后 loss 输出nan模型权重全部变成非法值。原因最常见的是学习率过大导致梯度爆炸其次是layer_norm的 eps 设置过小。在PositionwiseFeedForward里默认eps1e-6如果数据里有极端 outlier很容易溢出。解决把学习率降到5e-5以下同时在构建 model 时显式给eps1e-5如果依然 NaN检查数据是否包含空白行——空行会导致 token 序列为 0乘以位置 embedding 后产生 0 向量输入到 attention 里会除零。5.2 加载 state_dict 时键名不匹配现象torch.load后load_state_dict报missing keys或unexpected keys多打印键名差异后才明白是 composition 方式不同。原因本项目 embedding 是三个独立类HuggingFace 将三个 embedding 合并成一个大矩阵键名层次不同。解决写一个键名映射函数把embedding.token.weight替换成bert.embeddings.word_embeddings.weight这种形式再把多余的 dummy key 删掉。如果只是做推理更快的办法是跳过加载 project 自身的 embedding直接随机初始化 embedding 层因为下游任务一般也要重新训练 embedding。5.3 position embedding 超出训练时最大长度现象输入序列长度超过在PositionalEmbedding中设置的max_len512默认forward 时pe[:x.size(0)]切片产生越界错误。原因模型内部 buffer 只有 512 行位置编码向量输入 seq_len 512 时无位置编码可用。解决训练和推理时都把序列截断到 512如果需要更长序列把max_len改大然后重新初始化位置 embedding。注意这里不像有些项目能直接外推做长度扩展它没有做插值。5.4 普通公司 GPU 显存不够跑 batch_size 32现象我一度想用batch_size32、seq_len128复现论文设置结果显存直接爆掉。原因这个项目没有梯度累积机制一个 batch 就会撑爆显存。解决我习惯在 DataLoader 外面包一个梯度累积 loop每accum_steps4步做一次optimizer.step()相当于把 batch_size 等效放大。这个方法对这个项目完全适用只需在外面改几句代码不必动模型。5.5 中文语料的 token 切分异常现象直接喂一整段中文抽出来的 token 把每个字中间用空格断开结果词表爆炸模型学得很慢。原因默认的 tokenizer 用\w正则匹配中文会被当成连续的字串合在一起但如果语料里有空格空格也被当成分隔符结果碎片化。解决写一个简单逐字切分函数替代 tokenizer把句子按单字或按 jieba 分词后加入 vocab。只要保证 vocab 里覆盖全部 token模型就能正常训练只是过程中需要多花点心思设计自己的词表。6. 二次开发实践基于这个骨架做下游文本分类与序列标注6.1 改造思路替换输出头复用预训练特征这个 BERT-pytorch 最适合的二次开发方向是下游任务微调。通用做法是冻结 transformer 层参数只训练新增的任务头。拿文本分类举例改造分为三步加载BertModel后可取最后一层的[CLS]向量、或取最后一层所有 token 向量做池化之后接一个线性分类层再微调少量 epoch。class BertClassifier(nn.Module): def __init__(self, bert_model, num_classes): super(BertClassifier, self).__init__() self.bert bert_model self.classifier nn.Linear(self.bert.d_model, num_classes) def forward(self, tokens, segments): attn_mask (tokens ! 0).unsqueeze(1).unsqueeze(2) output, _ self.bert(tokens, segments, attn_mask) cls_hidden output[:, 0, :] logits self.classifier(cls_hidden) return logitstokens ! 0生成一个[batch, 1, seq_len, seq_len]的注意力 mask表示非 padding 位置参与计算。output[:, 0, :]取[CLS]对应的特征这是 BERT 分类任务的标准操作。此时bert可以是随机初始化的模型也可以加载此前训练出来的权重做迁移学习。6.2 验证方法和训练细节别直接上整套语料建议的做法是先拿小数据验证整个 loop 通不通。比如把seq_len32、epochs1跑一个只有 1000 条样本的小亿验证集确认 loss 在下降、分类 acc 能超过随机值后再放大规模。训练时我一般这样设置optimizer 用 AdamWlr2e-5warmup 比例 0.1通常两三个 epoch 就能达到不错效果。如果用全部参数微调对抗过拟合要配合dropout0.1和 attention dropout如果只训练分类层transformer 权重冻结需要较大的 lr如1e-3到5e-3。序列标注任务只需把num_classes改为标签类别数量输出维度从[batch, seq_len]变成[batch, seq_len, num_labels]然后对每个 token 位置算交叉熵注意 label 中 padding 位置要置为-100交叉熵忽略不要用 0 作为 ignore 系数因为真实标签编号可能包含 0。6.3 另外一个实用小技巧直接用这个框架写一个 mini 版 GPT如果你已经看懂了这个项目的 transformer 编排方式可以让Transformer模块只保留 decoder 侧的因式分解 mask 并换成标准的 GPT 式 causal maskdef bool_mask_from_seq_len(seq_len): mask torch.ones(seq_len, seq_len, dtypetorch.uint8).tril() return mask.unsqueeze(0)这个 mask 能保证对角线右侧全为 0使每个位置只能看到历史和当前 token实现自回归。改完之后别忘掉两个地方MultiHeadAttention的 masked_fill 中 mask0 会填 -1e9以及 drop attention 前要重新 Check 维度。这个简单的改造在学术实验里非常实用能让你用一个代码库横跨 BERT 和 GPT 两个家族。从那以后我每次拿到一个新的开源 NLP 项目都会强制走一遍同样的流程看model/目录下的类定义画一个结构图接着用最小语料跑通pretrain.py最后再把 attention mask 换成自己任务需要的形态。这样三步下来再复杂的预训练模型也能在两个小时内变成可用的下游工具。希望这个项目解析对你的 BERT 学习或二次开发能有一点实质帮助。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
