《动手学深度学习》word2vec 预训练数据集构建全流程：从 PTB 原始语料到可迭代小批量

人工智能深度学习机器学习教程【免费下载链接】d2l-zh《动手学深度学习》面向中文读者、能运行、可讨论。中英文版被70多个国家的500多所大学用于教学。项目地址https://gitcode.com/GitHub_Trending/d2/d2l-zh点击查看免费下载本文以 d2l-zh 仓库中 用于预训练词嵌入的数据集 一节为骨架英文原版见 word-embedding-dataset_origin.md完整讲解跳元模型skip-gram 负采样negative sampling训练前的数据准备流水线读取 PTB 语料、构建词表、高频词下采样、抽取中心词与上下文词、按 0.75 幂分布采样噪声词并最终通过batchify将变长样本整理成带掩码mask与标签label的定长小批量。读完本篇你将能独立复现load_data_ptb的完整实现并理解其中每个设计unk索引、下采样概率公式、mask/label 的一一对应关系背后的动机。文中所有函数均与仓库源码 d2l/torch.py 及 d2l/mxnet.py 中的实现一一对应。为什么需要专门的数据预处理阶段word2vec 的训练对象并不是原始的文本文件而是“中心词—上下文词”以及“中心词—噪声词”的配对关系。在讲解模型与近似训练方法见 跳元模型 与 近似训练之后本节解决的是从数据到训练样本的转换问题原始格式文本文件每行是一个句子单词以空格分隔中间格式过滤低频词、下采样高频词、抽取中心词及其上下文词、为每个上下文词采样噪声词最终格式按批量大小组织成可在训练中迭代加载的张量小批量并附带masks区分填充与非填充和labels区分正例与负例。读取数据集PTB 语料与词表构建数据集来源与注册本小节使用的语料是 Penn Tree BankPTB内容采样自《华尔街日报》文章被划分为训练集、验证集与测试集。原始格式中文本文件的每一行表示一个由空格分隔单词的句子这里把每个单词视为一个词元token。在 d2l 工具库中数据集通过DATA_HUB注册表管理注册项由“下载 URL SHA-1 校验值”组成下载时会先校验本地缓存文件的 SHA-1命中缓存则直接复用否则才发起下载见 d2l/torch.py 中的注册以及 d2l/torch.py 中download/download_extract的实现。PTB 的注册代码如下d2l.DATA_HUB[ptb] (d2l.DATA_URL ptb.zip, 319d85e578af0cdc590547f26231e4e31cdf1e42)read_ptb把文本行读成词元列表read_ptb负责下载并解压 PTB 训练集ptb.train.txt把整个文件按行拆分每行再按空白切分成词元列表def read_ptb(): Load the PTB dataset into a list of text lines. data_dir d2l.download_extract(ptb) # Read the training set. with open(os.path.join(data_dir, ptb.train.txt)) as f: raw_text f.read() return [line.split() for line in raw_text.split(\n)]调用后得到sentences——一个二维列表其长度即语料中的句子数量。构建词表min_freq10 与读取训练集后用d2l.Vocab(sentences, min_freq10)为整个语料构建词表任何出现次数少于 10 次的单词都会被替换为unk词元。需要注意PTB 原始数据中本身也包含表示稀有未知单词的unk词元。从源码 d2l/torch.py 可以看到Vocab类的实现细节通过count_corpusd2l/torch.py自动展平 1D/2D 词元列表后统计频率统计词频并按频率降序排序idx_to_token [unk] reserved_tokens即未知词元固定占索引 0unk属性返回 0遍历排序后的词频表当频率低于min_freq时停止加入词表从而实现“出现少于 10 次的词全部归入unk”__getitem__对不在词表中的词返回vocab.unk索引 0to_tokens则完成索引到词元的反向映射。词表构建完成后即可统计词表大小vocab d2l.Vocab(sentences, min_freq10) fvocab size: {len(vocab)}下采样以概率丢弃高频词动机文本数据中the、a、in这类高频词在超大语料中可能出现数十亿次但它们往往与上下文窗口中大量不同的词共现能提供的有用信号很少。例如上下文窗口中的chip它与低频词intel的共现显然比与高频词a的共现对训练更有价值。此外海量高频词的训练速度也慢。因此训练词嵌入模型时可以对高频词做下采样subsampling参见 Mikolov 等人 2013 年的 word2vec 论文。丢弃概率公式数据集中的每个词 $w_i$ 以如下概率被丢弃$$ P(w_i) \max\left(1 - \sqrt{\frac{t}{f(w_i)}}, 0\right),$$其中 $f(w_i)$ 是词 $w_i$ 的出现次数占数据集总词数的比例常量 $t$ 是超参数本节实验中取 $t 10^{-4}$。可以看到只有当相对频率满足 $f(w_i) t$ 时高频词 $w_i$ 才可能被丢弃且词的相对频率越高被丢弃的概率越大。subsample 实现def subsample(sentences, vocab): Subsample high-frequency words. # Exclude unknown tokens unk sentences [[token for token in line if vocab[token] ! vocab.unk] for line in sentences] counter d2l.count_corpus(sentences) num_tokens sum(counter.values()) # Return True if token is kept during subsampling def keep(token): return(random.uniform(0, 1) math.sqrt(1e-4 / counter[token] * num_tokens)) return ([[token for token in line if keep(token)] for line in sentences], counter)实现要点先剔除unk即索引为 0 的词元再统计词频与总词数keep函数按公式计算保留概率——从源码可见其等价于保留概率 $1 - P(w_i)\sqrt{\frac{10^{-4}}{f(w_i)}}\sqrt{\frac{t \cdot N}{count(w_i)}}$$N$ 为总词数。返回的第二个值counter是词频表供后续负采样权重计算复用。效果可视化与抽样率对比原文档用d2l.show_list_len_pair_hist实现在 d2l/torch.py同时绘制两组列表长度直方图对比下采样前后每句词元数量分布d2l.show_list_len_pair_hist([origin, subsampled], # tokens per sentence, count, sentences, subsampled);正如预期下采样通过删除高频词显著缩短句子从而加速训练。对单个词元做定量对比def compare_counts(token): return (f# of {token}: fbefore{sum([l.count(token) for l in sentences])}, fafter{sum([l.count(token) for l in subsampled])})高频词the采样率不到 1/20训练集中绝大多数the被丢弃低频词join被完全保留。下采样完成后把词元映射为索引得到可直接使用的数值型语料corpus [vocab[line] for line in subsampled] corpus[:3]提取中心词和上下文词get_centers_and_contexts 的逻辑get_centers_and_contexts从corpus中提取跳元模型训练所需的全部“中心词—上下文词”配对。它会对每条句子若句子长度不足 2无法构成中心词-上下文词对跳过对句子中每个位置 $i$用random.randint(1, max_window_size)均匀采样一个整数作为本次上下文窗口大小这样每个中心词的上下文范围在 1 到max_window_size之间随机变化取距离 $i$ 不超过窗口大小的索引范围range(max(0, i - window_size), min(len(line), i 1 window_size))并从中移除中心词自身 $i$剩余索引对应的词即为该中心词的上下文词。def get_centers_and_contexts(corpus, max_window_size): Return center words and context words in skip-gram. centers, contexts [], [] for line in corpus: # To form a center word--context word pair, each sentence needs to # have at least 2 words if len(line) 2: continue centers line for i in range(len(line)): # Context window centered at i window_size random.randint(1, max_window_size) indices list(range(max(0, i - window_size), min(len(line), i 1 window_size))) # Exclude the center word from the context words indices.remove(i) contexts.append([line[idx] for idx in indices]) return centers, contexts注意centers line意味着每个词都会作为一次中心词出现其对应的上下文词列表与自身位置一一对应因此centers与contexts长度相等。小数据集验证为直观验证逻辑构造一个包含 7 词与 3 词两个句子的微型数据集令最大上下文窗口为 2tiny_dataset [list(range(7)), list(range(7, 10))] print(dataset, tiny_dataset) for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)): print(center, center, has contexts, context)输出中可观察到句子首词只有右侧上下文末词只有左侧上下文且每个中心词的上下文数量随随机窗口大小浮动。在 PTB 上抽取全部配对在 PTB 数据集上训练时设置最大上下文窗口大小为 5一次性提取全部中心词与上下文词all_centers, all_contexts get_centers_and_contexts(corpus, 5) f# center-context pairs: {sum([len(contexts) for contexts in all_contexts])}该值即为跳元模型的“正样本”总数也是后续每个 epoch 要消费的样本量。负采样按 0.75 幂分布抽取噪声词RandomGenerator带缓存的加权随机抽取负采样需要根据预定义分布抽取噪声词。RandomGenerator类接收可能是未归一化的采样权重sampling_weights在 ${1, \dots, n}$ 上按权重随机抽取class RandomGenerator: Randomly draw among {1, ..., n} according to n sampling weights. def __init__(self, sampling_weights): # Exclude self.population list(range(1, len(sampling_weights) 1)) self.sampling_weights sampling_weights self.candidates [] self.i 0 def draw(self): if self.i len(self.candidates): # Cache k random sampling results self.candidates random.choices( self.population, self.sampling_weights, k10000) self.i 0 self.i 1 return self.candidates[self.i - 1]其核心优化是批量预采样缓存一次性用random.choices生成 10000 个候选draw时逐个取出缓存耗尽才重新采样。这一设计避免了高频调用random.choices的开销显著提升数据加载速度这也是本节练习 2 的关注点调整缓存大小k观察加载速度变化。用一组简单权重验证抽样分布在索引 1、2、3 上按概率 $P(X1)2/9, P(X2)3/9, P(X3)4/9$ 抽取generator RandomGenerator([2, 3, 4]) [generator.draw() for _ in range(10)]噪声词采样分布相对频率的 0.75 次幂根据 word2vec 论文的建议噪声词 $w$ 的采样概率 $P(w)$ 设为其在词典中相对频率的 0.75 次幂。对“中心词—上下文词”对随机抽取 $K$ 个实验中 $K5$噪声词。get_negatives的实现如下def get_negatives(all_contexts, vocab, counter, K): Return noise words in negative sampling. # Sampling weights for words with indices 1, 2, ... (index 0 is the # excluded unknown token) in the vocabulary sampling_weights [counter[vocab.to_tokens(i)]**0.75 for i in range(1, len(vocab))] all_negatives, generator [], RandomGenerator(sampling_weights) for contexts in all_contexts: negatives [] while len(negatives) len(contexts) * K: neg generator.draw() # Noise words cannot be context words if neg not in contexts: negatives.append(neg) all_negatives.append(negatives) return all_negatives两个关键细节索引 0unk被排除sampling_weights只覆盖索引 1 到len(vocab)-1对应源码注释“index 0 is the excluded unknown token”因为未知词元没有语义价值不应作为噪声样本噪声词不能是上下文词采样到的索引若出现在当前上下文词列表中则丢弃重抽避免“正样本被当成负样本”的冲突。对全部上下文词调用get_negatives(all_contexts, vocab, counter, 5)后all_negatives与all_contexts形状一一对应每个上下文词都配套了 $K$ 个噪声词索引。小批量加载训练实例batchify 与 mask/label变长样本的定长化问题提取完所有中心词、上下文词与噪声词后需要把它们组织成训练时可迭代加载的小批量。小批量中第 $i$ 个样本包括一个中心词、$n_i$ 个上下文词与 $m_i$ 个噪声词。由于上下文窗口随机变化$n_i m_i$ 随 $i$ 不同而变化。解决办法对每个样本把上下文词与噪声词拼接进contexts_negatives再用 0 填充到本批量内最大长度 $\max_i (n_im_i)$记为max_len。为了在损失计算中排除填充引入掩码变量masksmasks与contexts_negatives元素一一对应填充位置的masks取 0真实词元位置取 1。同时为区分正例上下文词与负例噪声词引入labels变量labels同样与contexts_negatives一一对应上下文词位置取 1噪声词位置取 0。batchify 实现def batchify(data): Return a minibatch of examples for skip-gram with negative sampling. max_len max(len(c) len(n) for _, c, n in data) centers, contexts_negatives, masks, labels [], [], [], [] for center, context, negative in data: cur_len len(context) len(negative) centers [center] contexts_negatives [context negative [0] * (max_len - cur_len)] masks [[1] * cur_len [0] * (max_len - cur_len)] labels [[1] * len(context) [0] * (max_len - len(context))] return (d2l.reshape(d2l.tensor(centers), (-1, 1)), d2l.tensor( contexts_negatives), d2l.tensor(masks), d2l.tensor(labels))返回的四个张量分别是centers形状 $(batch_size, 1)$reshape(-1, 1)把一维列表变成列向量、contexts_negatives形状 $(batch_size, max_len)$、masks与labels均与contexts_negatives同形状。下游训练时masks用于在损失求和时屏蔽填充位置labels用于告诉模型哪些是正例上下文词、哪些是负例噪声词。用两个样本验证x_1 (1, [2, 2], [3, 3, 3, 3]) x_2 (1, [2, 2, 2], [3, 3]) batch batchify((x_1, x_2)) names [centers, contexts_negatives, masks, labels] for name, data in zip(names, batch): print(name, , data)本例中$x_1$ 的上下文噪声长度为 6$x_2$ 的为 5故max_len6$x_2$ 末尾会多一个 0 填充对应地 $x_2$ 的masks末尾为 0$x_1$ 的labels为前 2 个 1上下文词 后 4 个 0$x_2$ 的labels为前 3 个 1 后 2 个 0。四者一一对应关系一目了然。整合代码load_data_ptb最后定义load_data_ptb把上述五个步骤串成完整流水线读取 PTB → 构建词表 → 下采样 → 抽取中心词/上下文词 → 负采样 → 打包为可迭代的数据加载器。仓库中为每个框架提供了独立实现对应原文档中#tab分块。MXNet 版本MXNet 版使用gluon.data.ArrayDataset包装三个等长列表再交给gluon.data.DataLoader并通过batchify_fnbatchify指定批量整理函数def load_data_ptb(batch_size, max_window_size, num_noise_words): Download the PTB dataset and then load it into memory. sentences read_ptb() vocab d2l.Vocab(sentences, min_freq10) subsampled, counter subsample(sentences, vocab) corpus [vocab[line] for line in subsampled] all_centers, all_contexts get_centers_and_contexts( corpus, max_window_size) all_negatives get_negatives( all_contexts, vocab, counter, num_noise_words) dataset gluon.data.ArrayDataset( all_centers, all_contexts, all_negatives) data_iter gluon.data.DataLoader( dataset, batch_size, shuffleTrue, batchify_fnbatchify, num_workersd2l.get_dataloader_workers()) return data_iter, vocabPyTorch 版本PyTorch 版定义了一个PTBDataset(torch.utils.data.Dataset)通过assert保证三个列表长度一致__getitem__按索引返回(center, context, negative)三元组torch.utils.data.DataLoader的collate_fnbatchify负责把一列三元组整理成定长小批量def load_data_ptb(batch_size, max_window_size, num_noise_words): Download the PTB dataset and then load it into memory. num_workers d2l.get_dataloader_workers() sentences read_ptb() vocab d2l.Vocab(sentences, min_freq10) subsampled, counter subsample(sentences, vocab) corpus [vocab[line] for line in subsampled] all_centers, all_contexts get_centers_and_contexts( corpus, max_window_size) all_negatives get_negatives( all_contexts, vocab, counter, num_noise_words) class PTBDataset(torch.utils.data.Dataset): def __init__(self, centers, contexts, negatives): assert len(centers) len(contexts) len(negatives) self.centers centers self.contexts contexts self.negatives negatives def __getitem__(self, index): return (self.centers[index], self.contexts[index], self.negatives[index]) def __len__(self): return len(self.centers) dataset PTBDataset(all_centers, all_contexts, all_negatives) data_iter torch.utils.data.DataLoader(dataset, batch_size, shuffleTrue, collate_fnbatchify, num_workersnum_workers) return data_iter, vocab两个版本的num_workers均来自d2l.get_dataloader_workers()其实现固定返回 4 个进程见 d2l/torch.py用于并行读取数据。调用与输出形状以批量大小 512、最大窗口 5、噪声词数 5 调用打印第一个小批量各变量的形状data_iter, vocab load_data_ptb(512, 5, 5) for batch in data_iter: for name, data in zip(names, batch): print(name, shape:, data.shape) break输出示意centers形状为(512, 1)contexts_negatives、masks、labels形状均为(512, max_len)其中max_len为本批内上下文噪声长度的最大值。返回的data_iter可直接用于 word2vec 预训练 一节中的模型训练循环vocab则用于训练结束后把学到的向量映射回词元。小结高频词在训练中提供的有用信号有限可以通过下采样丢弃它们来加速训练丢弃概率由 $P(w_i)\max(1-\sqrt{t/f(w_i)},0)$ 控制实验取 $t10^{-4}$。为提高计算效率样本以小批量方式加载由于上下文窗口随机变化导致样本长度不一需要用零填充定长化并额外定义变量来区分填充与非填充masks、正例与负例labels。整个流水线读数据 → 建词表 → 下采样 → 抽中心/上下文词 → 负采样 →batchify被封装进load_data_ptbMXNet 与 PyTorch 两套实现均可直接复用后续所有 word2vec 训练代码都依赖该数据接口。练习如果不使用下采样本节代码的运行时间会发生什么变化提示可把subsample返回的subsampled直接替换为sentences对比训练数据量。RandomGenerator类缓存k当前为 10000个随机采样结果。把k改为其他值如 100、1000、100000观察它对数据加载速度的影响分析缓存大小与内存、调用频率之间的权衡。本节代码中还有哪些超参数可能影响数据加载速度例如batch_size、max_window_size、噪声词数量K即num_noise_words、num_workers试分别改变它们并观察数据迭代与整体训练开销的变化。赞分享人工智能深度学习机器学习教程【免费下载链接】d2l-zh《动手学深度学习》面向中文读者、能运行、可讨论。中英文版被70多个国家的500多所大学用于教学。项目地址https://gitcode.com/GitHub_Trending/d2/d2l-zh点击查看免费下载相关推荐动手学深度学习word2vec 预训练词嵌入数据集PTB 下采样、负采样与小批量构建动手学深度学习word2vec 预训练词嵌入数据集PTB 下采样、负采样与小批量构建 导读 本文基于《动手学深度学习》d2l zh开源仓库中「用于预训人工智能深度学习机器学习教程d2l-en 词嵌入预训练数据集构建实战从 PTB 原始语料到 skip-gram 负采样小批量d2l en 词嵌入预训练数据集构建实战从 PTB 原始语料到 skip gram 负采样小批量 本文围绕《Dive into Deep Learning》文档教程人工智能深度学习NLP计算机视觉强化学习《动手学深度学习》实战用负采样在 PTB 语料上预训练 word2vec 跳元模型《动手学深度学习》实战用负采样在 PTB 语料上预训练 word2vec 跳元模型 本章承接《动手学深度学习》中 word2vec 理论基础 https://人工智能深度学习机器学习教程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考