- 人工智能
- 深度学习
- 机器学习
- 教程
【免费下载链接】d2l-zh
《动手学深度学习》：面向中文读者、能运行、可讨论。中英文版被70多个国家的500多所大学用于教学。
导读
本文讲解《动手学深度学习》中文版中 BERT 预训练数据集的完整构建流程：如何下载 WikiText-2 语料库，并将其转换为同时服务于**遮蔽语言模型（Masked Language Model, MLM）与下一句预测（Next Sentence Prediction, NSP）**两个预训练任务的标准样本格式。读完本文，你将掌握_read_wiki、_get_next_sentence、_replace_mlm_tokens、_pad_bert_inputs等核心函数的逐行原理，理解 15% 掩码比例、80/10/10 替换策略、序列填充与词表构建等关键设计，并能在 MXNet、PyTorch、PaddlePaddle 三个框架下直接复现训练数据管线。
本文以 chapter_natural-language-processing-pretraining/bert-dataset.md 为主体，其实现同时收录于仓库的 d2l/mxnet.py、d2l/torch.py、d2l/paddle.py 三个框架模块中，可与 chapter_natural-language-processing-pretraining/bert.md 中定义的 BERT 模型配套使用。
为什么需要自建 BERT 预训练数据集
最初的 BERT 模型在两个庞大的图书语料库与英语维基百科的合集上完成预训练，其数据规模与获取成本对大多数学习者并不友好。与此同时，现成的预训练 BERT 模型在医学等垂直领域往往表现不佳——领域词汇、句式与通用语料存在显著差异。因此在定制数据集上对 BERT 进行重新预训练已逐渐成为落地实践中的主流做法，这也是本文数据管线存在的根本动机。
为了让演示具备可复现性，原文档选择了体量适中、公开可得的WikiText-2语料库（论文引用自 Merity 等人 2016 年的工作）。与《动手学深度学习》中用于预训练 word2vec 的 PTB 数据集（见 chapter_natural-language-processing-pretraining/word2vec.md）相比，WikiText-2 具备三个关键优势：
- 保留了原始标点符号——这一点使其天然适配下一句预测任务，因为句号等标点是句子边界划分的基础；
- 保留了原始大小写与数字——信息熵更高，能训练出更贴合真实文本的表示；
- 语料规模大了一倍以上——在过滤低频词后，其词表量级仍远超 PTB（原文档运行示例中可直接打印
len(vocab)验证）。
数据集注册、下载与段落读取
在 DATA_HUB 中注册数据源
仓库的 d2l 工具模块维护了一个全局数据集注册表DATA_HUB。原文档将 WikiText-2 的下载地址与 SHA-1 校验值注册其中：
d2l.DATA_HUB['wikitext-2'] = ( 'https://s3.amazonaws.com/research.metamind.io/wikitext/' 'wikitext-2-v1.zip', '3c914d17d80b1459be871a5039ac23e752a53cbe')在 d2l/torch.py 等框架模块中保留了完全一致的注册代码。d2l.download_extract会依据注册信息完成「下载 → SHA-1 校验 → 解压」三步操作（见 d2l/torch.py），其中 SHA-1 校验用于确保下载文件完整性，命中本地缓存后不会重复下载。
段落级读取_read_wiki
WikiText-2 的原始文件wiki.train.tokens中每一行代表一个段落，段落内任意标点符号与其前面的词元之间都插入了空格。_read_wiki的职责是从原始文件还原出结构化的段落列表：
def _read_wiki(data_dir): file_name = os.path.join(data_dir, 'wiki.train.tokens') with open(file_name, 'r') as f: lines = f.readlines() # 大写字母转换为小写字母 paragraphs = [line.strip().lower().split(' . ') for line in lines if len(line.split(' . ')) >= 2] random.shuffle(paragraphs) return paragraphs关键细节逐条拆解：
- 只保留至少包含两句话的段落：由于下一句预测需要「句子对」，
len(line.split(' . ')) >= 2的过滤条件保证了每个保留段落都能产出至少一个 NSP 训练样本； - 仅以句号
' . '作为句子分隔符：这是文档明确声明的简化处理——句子拆分的通用做法（如 Spacy、NLTK）会留在文末练习中讨论； - 统一转为小写：降低词表规模、简化匹配，代价是丢失大小写信息；
- 随机打乱段落顺序：避免相邻训练样本之间的语料偏置，等价于在数据层面做了一次 shuffle。
为预训练任务定义辅助函数
BERT 的预训练同时包含两个自监督任务，原始文本必须先转换成各自的样本格式。以下辅助函数正是为此设计，最终统一由_WikiTextDataset组装调用。
生成下一句预测（NSP）任务的数据
NSP 是一个二分类任务：给定句子 A 与句子 B，判断 B 是否为 A 在原文中的下一句。_get_next_sentence负责构造单个句子对样本：
def _get_next_sentence(sentence, next_sentence, paragraphs): if random.random() < 0.5: is_next = True else: # paragraphs是三重列表的嵌套 next_sentence = random.choice(random.choice(paragraphs)) is_next = False return sentence, next_sentence, is_next它以 50% 概率保留真实的相邻句子（标签is_next = True），以另外 50% 概率从整个语料库的任意段落中随机抽取一个句子替换句子 B（标签is_next = False）。这里paragraphs是「段落 → 句子 → 词元」的三重嵌套列表，random.choice(random.choice(paragraphs))先随机选段落、再随机选句子。
_get_nsp_data_from_paragraph则把单个段落中所有相邻句子对依次喂给上述函数，生成该段落的全部 NSP 样本：
def _get_nsp_data_from_paragraph(paragraph, paragraphs, vocab, max_len): nsp_data_from_paragraph = [] for i in range(len(paragraph) - 1): tokens_a, tokens_b, is_next = _get_next_sentence( paragraph[i], paragraph[i + 1], paragraphs) # 考虑1个'<cls>'词元和2个'<sep>'词元 if len(tokens_a) + len(tokens_b) + 3 > max_len: continue tokens, segments = d2l.get_tokens_and_segments(tokens_a, tokens_b) nsp_data_from_paragraph.append((tokens, segments, is_next)) return nsp_data_from_paragraph注意其中的max_len长度约束：BERT 输入序列包含 1 个<cls>词元和 2 个<sep>词元，因此只有当len(tokens_a) + len(tokens_b) + 3不超过最大序列长度时才保留该样本，超长句子对直接丢弃。d2l.get_tokens_and_segments（实现见 d2l/torch.py）负责把两个句子包装成标准 BERT 输入：['<cls>'] + tokens_a + ['<sep>'] + tokens_b + ['<sep>']，并用 0/1 分别标记片段 A 与片段 B。
生成遮蔽语言模型（MLM）任务的数据
MLM 要求模型从被破坏的输入中预测被掩蔽词元的真实标签。_replace_mlm_tokens实现词元替换策略：
def _replace_mlm_tokens(tokens, candidate_pred_positions, num_mlm_preds, vocab): # 为遮蔽语言模型的输入创建新的词元副本，其中输入可能包含替换的“<mask>”或随机词元 mlm_input_tokens = [token for token in tokens] pred_positions_and_labels = [] # 打乱后用于在遮蔽语言模型任务中获取15%的随机词元进行预测 random.shuffle(candidate_pred_positions) for mlm_pred_position in candidate_pred_positions: if len(pred_positions_and_labels) >= num_mlm_preds: break masked_token = None # 80%的时间：将词替换为“<mask>”词元 if random.random() < 0.8: masked_token = '<mask>' else: # 10%的时间：保持词不变 if random.random() < 0.5: masked_token = tokens[mlm_pred_position] # 10%的时间：用随机词替换该词 else: masked_token = random.choice(vocab.idx_to_token) mlm_input_tokens[mlm_pred_position] = masked_token pred_positions_and_labels.append( (mlm_pred_position, tokens[mlm_pred_position])) return mlm_input_tokens, pred_positions_and_labels这段代码是 BERT 原始论文掩码策略的忠实复刻，其概率分布如下：
| 情形 | 概率 | 输入中的处理 |
|---|---|---|
替换为<mask>特殊词元 | 80% | mlm_input_tokens[pos] = '<mask>' |
| 保持原词不变 | 10% | mlm_input_tokens[pos] = tokens[pos] |
| 替换为词表中随机词元 | 10% | mlm_input_tokens[pos] = random.choice(vocab.idx_to_token) |
之所以不 100% 使用<mask>，是为了缓解预训练与下游微调之间的「输入分布不匹配」：微调阶段输入中并不存在<mask>词元。保留原词与注入随机词元的少量噪声，能鼓励 BERT 的双向上下文编码不过度依赖掩码标记本身。
_get_mlm_data_from_tokens负责确定预测哪些位置：
def _get_mlm_data_from_tokens(tokens, vocab): candidate_pred_positions = [] # tokens是一个字符串列表 for i, token in enumerate(tokens): # 在遮蔽语言模型任务中不会预测特殊词元 if token in ['<cls>', '<sep>']: continue candidate_pred_positions.append(i) # 遮蔽语言模型任务中预测15%的随机词元 num_mlm_preds = max(1, round(len(tokens) * 0.15)) mlm_input_tokens, pred_positions_and_labels = _replace_mlm_tokens( tokens, candidate_pred_positions, num_mlm_preds, vocab) pred_positions_and_labels = sorted(pred_positions_and_labels, key=lambda x: x[0]) pred_positions = [v[0] for v in pred_positions_and_labels] mlm_pred_labels = [v[1] for v in pred_positions_and_labels] return vocab[mlm_input_tokens], pred_positions, vocab[mlm_pred_labels]三个设计要点：
- 特殊词元不参与预测：
<cls>与<sep>被显式排除出候选位置，因为它们承担的是序列级语义汇聚与句子边界标记职责； - 预测数量为序列长度的 15%（向上取整）：
num_mlm_preds = max(1, round(len(tokens) * 0.15))，max(1, ...)保证即使是极短序列也至少预测 1 个位置； - 输出统一为索引：函数最后通过
vocab[...]把字符串词元批量映射为索引，返回(输入词元索引, 预测位置, 预测标签索引)三元组，可直接送入交叉熵损失计算。
将文本转换为预训练数据集
统一填充_pad_bert_inputs
真实语料中句子长度参差不齐，深度学习框架要求 batch 内张量形状一致，因此需要填充。_pad_bert_inputs接受来自两个预训练任务辅助函数的输出样例，输出 7 个等长对齐的数组：
def _pad_bert_inputs(examples, max_len, vocab): max_num_mlm_preds = round(max_len * 0.15) all_token_ids, all_segments, valid_lens, = [], [], [] all_pred_positions, all_mlm_weights, all_mlm_labels = [], [], [] nsp_labels = [] for (token_ids, pred_positions, mlm_pred_label_ids, segments, is_next) in examples: all_token_ids.append(np.array(token_ids + [vocab['<pad>']] * ( max_len - len(token_ids)), dtype='int32')) all_segments.append(np.array(segments + [0] * ( max_len - len(segments)), dtype='int32')) # valid_lens不包括'<pad>'的计数 valid_lens.append(np.array(len(token_ids), dtype='float32')) all_pred_positions.append(np.array(pred_positions + [0] * ( max_num_mlm_preds - len(pred_positions)), dtype='int32')) # 填充词元的预测将通过乘以0权重在损失中过滤掉 all_mlm_weights.append( np.array([1.0] * len(mlm_pred_label_ids) + [0.0] * ( max_num_mlm_preds - len(pred_positions)), dtype='float32')) all_mlm_labels.append(np.array(mlm_pred_label_ids + [0] * ( max_num_mlm_preds - len(mlm_pred_label_ids)), dtype='int32')) nsp_labels.append(np.array(is_next)) return (all_token_ids, all_segments, valid_lens, all_pred_positions, all_mlm_weights, all_mlm_labels, nsp_labels)逐字段说明如下：
| 输出数组 | 含义 | 填充方式 |
|---|---|---|
all_token_ids | 输入序列的词元索引 | 尾部补<pad>到max_len |
all_segments | 片段 A/B 标记（0/1） | 尾部补 0 到max_len |
valid_lens | 有效长度（不含<pad>） | 记录真实长度 |
all_pred_positions | 被预测位置的下标 | 尾部补 0 到max_num_mlm_preds |
all_mlm_weights | 预测位置权重掩码 | 真实预测处为 1.0，填充处为 0.0 |
all_mlm_labels | MLM 预测标签索引 | 尾部补 0 到max_num_mlm_preds |
nsp_labels | NSP 二分类标签（0/1） | 无需填充 |
其中all_mlm_weights是实现填充位置「零贡献」的关键：预训练阶段计算 MLM 交叉熵损失时会用该权重逐位置相乘（见 d2l/torch.py 中的_get_batch_loss_bert，mlm_l = mlm_l.sum() / (mlm_weights_X.sum() + 1e-8)），填充位置的权重为 0，从而在损失中被完全过滤，避免无意义的填充槽位干扰梯度。
组装数据集_WikiTextDataset
_WikiTextDataset将前述所有辅助函数串联成完整流水线：
class _WikiTextDataset(torch.utils.data.Dataset): def __init__(self, paragraphs, max_len): # 输入paragraphs[i]是代表段落的句子字符串列表； # 而输出paragraphs[i]是代表段落的句子列表，其中每个句子都是词元列表 paragraphs = [d2l.tokenize( paragraph, token='word') for paragraph in paragraphs] sentences = [sentence for paragraph in paragraphs for sentence in paragraph] self.vocab = d2l.Vocab(sentences, min_freq=5, reserved_tokens=[ '<pad>', '<mask>', '<cls>', '<sep>']) # 获取下一句子预测任务的数据 examples = [] for paragraph in paragraphs: examples.extend(_get_nsp_data_from_paragraph( paragraph, paragraphs, self.vocab, max_len)) # 获取遮蔽语言模型任务的数据 examples = [(_get_mlm_data_from_tokens(tokens, self.vocab) + (segments, is_next)) for tokens, segments, is_next in examples] # 填充输入 (self.all_token_ids, self.all_segments, self.valid_lens, self.all_pred_positions, self.all_mlm_weights, self.all_mlm_labels, self.nsp_labels) = _pad_bert_inputs( examples, max_len, self.vocab) def __getitem__(self, idx): return (self.all_token_ids[idx], self.all_segments[idx], self.valid_lens[idx], self.all_pred_positions[idx], self.all_mlm_weights[idx], self.all_mlm_labels[idx], self.nsp_labels[idx]) def __len__(self): return len(self.all_token_ids)流水线的四个阶段清晰可辨：
- 词元化：
d2l.tokenize(paragraph, token='word')将句子字符串按空格切分为词元列表（实现见 d2l/torch.py，word 模式等价于line.split()）； - 构建词表：
d2l.Vocab(sentences, min_freq=5, reserved_tokens=['<pad>', '<mask>', '<cls>', '<sep>'])以全部句子为统计语料构建词表，出现次数少于 5 次的低频词元被过滤，四个特殊词元（含<unk>，见 d2l/torch.py）占据词表头部固定索引。原文档同时指出：最初的 BERT 使用词表大小为 30000 的 WordPiece 嵌入，其词元化是对字节对编码（BPE）算法的微调；为简化演示，此处使用 word 级切分； - 生成双任务样本：先为每个段落生成 NSP 样本，再对每个 NSP 样本的 token 序列追加 MLM 样本（通过
_get_mlm_data_from_tokens(...) + (segments, is_next)拼接三元组与片段/标签信息）； - 统一填充：调用
_pad_bert_inputs输出 7 个等长数组，并通过__getitem__支持按索引访问单条样本、通过__len__报告样本总数。
统一加载入口load_data_wiki
load_data_wiki封装了「下载 → 读取 → 建数据集 → 建 DataLoader」的完整流程：
def load_data_wiki(batch_size, max_len): """加载WikiText-2数据集""" num_workers = d2l.get_dataloader_workers() data_dir = d2l.download_extract('wikitext-2', 'wikitext-2') paragraphs = _read_wiki(data_dir) train_set = _WikiTextDataset(paragraphs, max_len) train_iter = torch.utils.data.DataLoader(train_set, batch_size, shuffle=True, num_workers=num_workers) return train_iter, train_set.vocab两个入参的意义：batch_size控制小批量大小；max_len为 BERT 输入序列最大长度，它同时决定 MLM 每样本最多预测round(max_len * 0.15)个位置。d2l.get_dataloader_workers()（见 d2l/torch.py）返回 4 个进程用于数据读取。函数返回训练迭代器与词表对象，后者将用于后续预训练阶段的嵌入层尺寸设定。
运行示例：小批量样本形状与词表规模
原文档给出了一组可直接运行的默认参数，将批量大小设为 512、最大序列长度设为 64：
batch_size, max_len = 512, 64 train_iter, vocab = load_data_wiki(batch_size, max_len) for (tokens_X, segments_X, valid_lens_x, pred_positions_X, mlm_weights_X, mlm_Y, nsp_y) in train_iter: print(tokens_X.shape, segments_X.shape, valid_lens_x.shape, pred_positions_X.shape, mlm_weights_X.shape, mlm_Y.shape, nsp_y.shape) break该循环只消费一个小批量用于观察形状。由于max_len = 64，每个输入序列中为 MLM 任务预测的位置数恰为round(64 × 0.15) = 10，因此pred_positions_X、mlm_weights_X、mlm_Y的最后一维均为 10；valid_lens_x记录每条样本去除<pad>后的真实长度；nsp_y保存批内 512 条 NSP 标签。
随后可验证词表规模：
len(vocab)原文档明确指出：即使过滤掉出现次数少于 5 次的低频词元，WikiText-2 构建出的词表仍比 PTB 数据集大两倍以上，这也是该语料信息量更丰富的直接证据。
三种框架的统一实现
本文数据管线并非某一框架专属：原文档及 d2l 模块为 MXNet、PyTorch、PaddlePaddle 分别维护了等价实现，差异仅体现在数据容器与张量 API 上：
| 框架 | 数据集基类 | DataLoader | 张量创建 | 实现位置 |
|---|---|---|---|---|
| MXNet | gluon.data.Dataset | gluon.data.DataLoader | np.array(..., dtype='int32') | d2l/mxnet.py（wikitext-2注册于第 2192 行起） |
| PyTorch | torch.utils.data.Dataset | torch.utils.data.DataLoader | torch.tensor(..., dtype=torch.long/float32) | d2l/torch.py |
| PaddlePaddle | paddle.io.Dataset | paddle.io.DataLoader | paddle.to_tensor(..., dtype=paddle.int64/float32) | d2l/paddle.py（wikitext-2注册于第 2331 行起） |
三者的函数名、参数语义与返回结构完全一致，_pad_bert_inputs与_WikiTextDataset只是随框架分别定义（原文档依次给出了 mxnet/pytorch/paddle 三个代码块）。这意味着同一套数据处理思路可直接迁移到任意框架，后续 chapter_natural-language-processing-pretraining/bert-pretraining.md 章节即可在该数据迭代器上驱动BERTModel完成预训练。
小结
- WikiText-2 保留了原始标点符号、大小写和数字，规模是 PTB 的两倍以上，适合作为演示级 BERT 预训练语料；
- NSP 样本以 50% 概率由真实相邻句子构成正例，以 50% 概率由随机句子构成负例；
- MLM 样本对 15% 的候选词元执行 80/10/10 替换（
<mask>/ 原词 / 随机词），特殊词元<cls>、<sep>不参与预测； _pad_bert_inputs通过mlm_weights权重掩码将填充位置的预测损失归零，_WikiTextDataset通过__getitem__支持任意索引访问句子对生成的预训练样本；- 整个管线在 MXNet、PyTorch、PaddlePaddle 三框架中实现一致，可通过
load_data_wiki(batch_size, max_len)一行代码加载。
延伸练习
- 升级句子拆分技术：原实现仅以句号
' . '为分隔符，可尝试 Spacy 或 NLTK 等更完善的句子分词器。以 NLTK 为例：先pip install nltk，代码中import nltk，再nltk.download('punkt')；对'This is great ! Why not ?'调用nltk.tokenize.sent_tokenize将返回['This is great !', 'Why not ?']两个句子。替换拆分逻辑后观察 NSP 正样本质量与整体样本数的变化。 - 观察词表膨胀效应：
Vocab的min_freq=5过滤了低频词元。若去掉该过滤条件（min_freq=0）重新构建词表，思考词量会膨胀到什么量级，以及这会对词嵌入矩阵的参数量与训练效率带来什么影响。
- 人工智能
- 深度学习
- 机器学习
- 教程
【免费下载链接】d2l-zh
《动手学深度学习》：面向中文读者、能运行、可讨论。中英文版被70多个国家的500多所大学用于教学。
相关推荐
《动手学深度学习》实战：基于 WikiText-2 从零预训练 BERT 模型（遮蔽语言模型 + 下一句预测）
《动手学深度学习》实战：基于 WikiText 2 从零预训练 BERT 模型（遮蔽语言模型 + 下一句预测） 本篇技术指南以 d2l zh https://l
人工智能深度学习机器学习教程动手学深度学习：用 WikiText-2 构建 BERT 预训练数据集（MLM 与 NSP 全流程解析）
动手学深度学习：用 WikiText 2 构建 BERT 预训练数据集（MLM 与 NSP 全流程解析） 在《动手学深度学习》（d2l zh）中，BERT 的预
人工智能深度学习机器学习教程BERT 预训练数据集构建实战：基于 d2l-en 的 WikiText-2 掩码语言建模与下一句预测数据管线
BERT 预训练数据集构建实战：基于 d2l en 的 WikiText 2 掩码语言建模与下一句预测数据管线 导读 本文以《动手学深度学习》（d2l en）的
文档教程人工智能深度学习NLP计算机视觉强化学习
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
