BERT-PyTorch 复现指南：Transformer 预训练模型结构全解析

简介：面向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 的机器上都能复现。

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 / B，position 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_idx=0)

这个类直接继承nn.Embedding，内部就是一个查表操作。padding_idx=0意味着词表中索引为 0 的 token（通常是[PAD]）在梯度更新时会被忽略，避免 padding 部分干扰训练。vocab_size在你的语料上训练 tokenizer 后确定，常见值像 30522（BERT-base 原始词表）。

Segment embedding 稍微不一样，它要区分两个句子。所以SegmentEmbedding是一个 embedding 矩阵，但 vocab_size 固定为 2——type 0 代表句子 A，type 1 代表句子 B：

class SegmentEmbedding(nn.Embedding): def __init__(self, embed_size): super(SegmentEmbedding, self).__init__(2, embed_size, padding_idx=0)

这里embed_size就是d_model（768）。输入给它的 tensor 是[batch, seq_len]的 0 / 1 序列，输出[batch, seq_len, 768]的段向量。注意padding_idx=0在这个类里表示 segment 标记为 0 的位置也是句子 A 的一部分，不是真的 padding，这不影响使用。

3.2 Positional Embedding：正弦函数位置编码的实现

BERT 没有 RNN / CNN 这种自带顺序的结构，所以必须显式注入位置信息。项目里的位置编码沿用了 transformer 论文里的正弦余弦函数方案：

class PositionalEmbedding(nn.Module): def __init__(self, d_model, max_len=512): super(PositionalEmbedding, self).__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.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, mask=None): 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, attn

forward 中先通过三个线性层得到 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会报错。

真正的注意力计算封装在ScaledDotProductAttention：

class ScaledDotProductAttention(nn.Module): def __init__(self, d_k, attn_dropout=0.1): super(ScaledDotProductAttention, self).__init__() self.dropout = nn.Dropout(attn_dropout) def forward(self, q, k, v, mask=None): 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, attn

masked_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, dropout=0.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, eps=1e-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 x

d_in是 768，d_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, dropout=0.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_mask=None): residual = enc_input enc_input, attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_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 维向量，过一个线性层映射回词表大小，再用交叉熵计算 loss：

masked_loss = self.criterion(masked_output, masked_label) # criterion = torch.nn.CrossEntropyLoss(ignore_index=0)

ignore_index=0表示 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

常用参数含义：

如果你是新手，第一次跑建议把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里默认eps=1e-6，如果数据里有极端 outlier，很容易溢出。

解决：把学习率降到5e-5以下；同时在构建 model 时显式给eps=1e-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_len=512（默认），forward 时pe[:x.size(0)]切片产生越界错误。

原因：模型内部 buffer 只有 512 行位置编码向量，输入 seq_len > 512 时无位置编码可用。

解决：训练和推理时都把序列截断到 512；如果需要更长序列，把max_len改大然后重新初始化位置 embedding。注意这里不像有些项目能直接外推做长度扩展，它没有做插值。

5.4 普通公司 GPU 显存不够跑 batch_size 32

现象：我一度想用batch_size=32、seq_len=128复现论文设置，结果显存直接爆掉。

原因：这个项目没有梯度累积机制，一个 batch 就会撑爆显存。

解决：我习惯在 DataLoader 外面包一个梯度累积 loop：每accum_steps=4步做一次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 logits

tokens != 0生成一个[batch, 1, seq_len, seq_len]的注意力 mask，表示非 padding 位置参与计算。output[:, 0, :]取[CLS]对应的特征，这是 BERT 分类任务的标准操作。此时bert可以是随机初始化的模型，也可以加载此前训练出来的权重做迁移学习。

6.2 验证方法和训练细节：别直接上整套语料

建议的做法是：先拿小数据验证整个 loop 通不通。比如把seq_len=32、epochs=1，跑一个只有 1000 条样本的小亿验证集，确认 loss 在下降、分类 acc 能超过随机值后，再放大规模。训练时我一般这样设置：

• optimizer 用 AdamW，lr=2e-5，warmup 比例 0.1，通常两三个 epoch 就能达到不错效果。
• 如果用全部参数微调，对抗过拟合要配合dropout=0.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 mask：

def bool_mask_from_seq_len(seq_len): mask = torch.ones(seq_len, seq_len, dtype=torch.uint8).tril() return mask.unsqueeze(0)

这个 mask 能保证对角线右侧全为 0，使每个位置只能看到历史和当前 token，实现自回归。改完之后别忘掉两个地方：MultiHeadAttention的 masked_fill 中 mask=0 会填 -1e9；以及 drop attention 前要重新 Check 维度。这个简单的改造在学术实验里非常实用，能让你用一个代码库横跨 BERT 和 GPT 两个家族。

从那以后，我每次拿到一个新的开源 NLP 项目，都会强制走一遍同样的流程：看model/目录下的类定义，画一个结构图；接着用最小语料跑通pretrain.py；最后再把 attention mask 换成自己任务需要的形态。这样三步下来，再复杂的预训练模型也能在两个小时内变成可用的下游工具。希望这个项目解析对你的 BERT 学习或二次开发能有一点实质帮助。

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