D2L 实战：从零预训练 word2vec 的 skip-gram 模型（负采样 + PTB 数据集）

• 文档
• 教程
• 人工智能
• 深度学习
• NLP
• 计算机视觉
• 强化学习

【免费下载链接】d2l-en

Interactive deep learning book with multi-framework code, math, and discussions. Adopted at 500 universities from 70 countries including Stanford, MIT, Harvard, and Cambridge.

项目地址：https://gitcode.com/gh_mirrors/d2/d2l-en

点击查看免费下载

本篇文章以 D2L（Dive into Deep Learning）仓库中chapter_natural-language-processing-pretraining/word2vec-pretraining.md为骨架，完整演示如何基于 PTB（Penn Tree Bank）语料，用 PyTorch / MXNet 双框架从零实现并预训练 word2vec 的 skip-gram 模型：包括数据加载、嵌入层与前向传播、带掩码的二元交叉熵损失、训练循环，以及训练完成后用余弦相似度检索语义相近词。读完本文，你将掌握一套可直接复制运行的 word2vec 预训练流水线，并理解其中 mask、label、负采样等关键设计的底层原理。

1. 整体路线：从数据迭代器到语义相似的词

预训练 word2vec 的完整链路可以分为四步：

1. 通过d2l.load_data_ptb获得 PTB 数据集的数据迭代器data_iter与词表vocab；
2. 用嵌入层（Embedding Layer）与批量矩阵乘法实现 skip-gram 模型的前向传播；
3. 使用负采样（negative sampling）对应的二元交叉熵损失训练模型（两个嵌入层分别对应中心词向量与上下文词向量）；
4. 训练完成后，取中心词嵌入层的权重，用余弦相似度找出与查询词最相似的词。

其中模型本身的数学定义（中心词预测上下文词的条件概率、softmax 形式与训练损失）来自 word2vec.md，数据集处理（下采样、中心词/上下文词抽取、负采样、minibatch 构造）来自 word-embedding-dataset.md，负采样与近似训练的理论来自 approx-training.md。本文聚焦“实现与预训练”这一环节，并穿插源码级讲解。

2. 准备数据：调用d2l.load_data_ptb

训练开始前，先获取数据迭代器和词表。本文的预训练实验使用以下超参数：批量大小 512、最大上下文窗口 5、每个（中心词，上下文词）对采样的噪声词数 5。

#@tab mxnet from d2l import mxnet as d2l import math from mxnet import autograd, gluon, np, npx from mxnet.gluon import nn npx.set_np() batch_size, max_window_size, num_noise_words = 512, 5, 5 data_iter, vocab = d2l.load_data_ptb(batch_size, max_window_size, num_noise_words)

#@tab pytorch from d2l import torch as d2l import math import torch from torch import nn batch_size, max_window_size, num_noise_words = 512, 5, 5 data_iter, vocab = d2l.load_data_ptb(batch_size, max_window_size, num_noise_words)

2.1 数据流水线的底层实现

load_data_ptb是仓库中真实可用的工具函数，PyTorch 版本定义在 d2l/torch.py，MXNet 版本定义在 d2l/mxnet.py。其内部按顺序完成如下环节：

1. 读取语料：read_ptb()从d2l.DATA_HUB['ptb']（PTB 压缩包，校验和为319d85e578af0cdc590547f26231e4e31cdf1e42）下载并解压数据，按行读取ptb.train.txt，每行按空格切分成一个句子。PTB 语料采样自《华尔街日报》文章，原始格式就是“每行一句、词与词以空格分隔”。
2. 构建词表：d2l.Vocab(sentences, min_freq=10)将出现次数低于 10 次的词统一替换为<unk>未知词元；原始数据中本就包含<unk>。
3. 高频词下采样：subsample()以概率 $P(w_i)=\max(1-\sqrt{t/f(w_i)},0)$ 丢弃高频词，其中 $f(w_i)$ 是该词的相对频率，阈值 $t=10^{-4}$（见 d2l/torch.py）。例如“the”的保留率不到 1/20，而低频词如 “join” 会被完整保留。下采样既能减少“the、a、in”这类与大量词共现、信息量低的高频词，也能显著加速训练。
4. 抽取中心词与上下文词：get_centers_and_contexts(corpus, max_window_size)对每个句子、每个位置i随机采样一个 1 到max_window_size之间的整数作为窗口大小，窗口内除中心词外的词即为上下文词（见 d2l/torch.py）。随机窗口大小等效于对窗口尺寸做了数据增强，能有效利用全部上下文信息。
5. 负采样：get_negatives(all_contexts, vocab, counter, K)为每个上下文词对采样 K 个噪声词（本文 K=5），噪声词按词频的 0.75 次幂加权分布抽样，且不能是当前的上下文词（见 d2l/torch.py）。抽样由RandomGenerator实现，它一次性缓存 10000 个抽样结果，避免频繁调用随机数生成器拖慢数据加载。
6. 构造 minibatch：batchify()将每个样本的上下文词与噪声词拼接为contexts_negatives，并对齐到批内最长长度max_len；同时生成掩码masks（1 表示真实词元，0 表示填充位）和标签labels（1 表示正例上下文词，0 表示噪声词与填充位），见 d2l/torch.py。返回的四个张量形状为：中心词(batch_size, 1)、contexts_negatives (batch_size, max_len)、masks (batch_size, max_len)、labels (batch_size, max_len)。

训练时每个 minibatch 的max_len由当前批内各样本的实际长度决定（max(len(c)+len(n) for _, c, n in data)），因此不同批次间长度可以不同。PyTorch 端通过DataLoader(..., collate_fn=batchify, num_workers=d2l.get_dataloader_workers())注入该批处理函数，MXNet 端则使用gluon.data.DataLoader(..., batchify_fn=batchify)，两者行为一致。

3. 实现 Skip-Gram 模型

skip-gram 模型的实现只用嵌入层与批量矩阵乘法两样东西即可完成，不需要任何复杂的神经网络结构。

3.1 回顾嵌入层（Embedding Layer）

嵌入层把词元的索引映射为特征向量：其权重是一个矩阵，行数等于字典大小（input_dim/num_embeddings），列数等于每个词元的向量维度（output_dim/embedding_dim）。词嵌入模型训练完成后，这个权重矩阵就是我们最终需要的词向量。

#@tab mxnet embed = nn.Embedding(input_dim=20, output_dim=4) embed.initialize() embed.weight

#@tab pytorch embed = nn.Embedding(num_embeddings=20, embedding_dim=4) print(f'Parameter embedding_weight ({embed.weight.shape}, ' f'dtype={embed.weight.dtype})')

嵌入层的输入是词元（单词）的索引：对任意索引 $i$，其向量表示就是权重矩阵的第 $i$ 行。向量维度设为 4 时，对形状为(2, 3)的 token 索引小批量，嵌入层返回形状(2, 3, 4)的向量：

#@tab all x = d2l.tensor([[1, 2, 3], [4, 5, 6]]) embed(x)

3.2 定义前向传播：skip_gram

前向传播的输入包括两部分：

• 中心词索引center，形状为(batch size, 1)；
• 拼接后的上下文词与噪声词索引contexts_and_negatives，形状为(batch size, max_len)，其中max_len由小批量加载阶段确定（见 word-embedding-dataset.md 的 minibatch 部分）。

这两个变量先分别通过嵌入层embed_v（中心词向量）与embed_u（上下文/噪声词向量）转成向量，再做批量矩阵乘法，输出形状为(batch size, 1, max_len)。输出中的每个元素是一个中心词向量与一个上下文（或噪声）词向量的点积：

#@tab mxnet def skip_gram(center, contexts_and_negatives, embed_v, embed_u): v = embed_v(center) u = embed_u(contexts_and_negatives) pred = npx.batch_dot(v, u.swapaxes(1, 2)) return pred

#@tab pytorch def skip_gram(center, contexts_and_negatives, embed_v, embed_u): v = embed_v(center) u = embed_u(contexts_and_negatives) pred = torch.bmm(v, u.permute(0, 2, 1)) return pred

这里批量矩阵乘法等价于把形状(batch, 1, embed_size)的v与转置后形状(batch, embed_size, max_len)的u相乘，得到(batch, 1, max_len)的相似度得分。打印几个示例输入下的输出形状：

#@tab mxnet skip_gram(np.ones((2, 1)), np.ones((2, 4)), embed, embed).shape

#@tab pytorch skip_gram(torch.ones((2, 1), dtype=torch.long), torch.ones((2, 4), dtype=torch.long), embed, embed).shape

4. 训练：负采样 + 二元交叉熵损失

4.1 从负采样目标到二元交叉熵损失

在 approx-training.md 中，负采样把“上下文词 $w_o$ 出现在中心词 $w_c$ 的上下文窗口内”建模为概率事件 $P(D=1\mid w_c, w_o)=\sigma(\mathbf{u}_o^\top \mathbf{v}_c)$。如果只最大化所有正例事件的联合概率，目标会在所有词向量趋于无穷时取到最大值 1，毫无意义；因此负采样还要加入从预定义分布中采样的负例，把目标函数变成区分正例与负例的二分类问题——这正是二元交叉熵损失（binary cross-entropy loss）的用武之地。

#@tab mxnet loss = gluon.loss.SigmoidBCELoss()

#@tab pytorch class SigmoidBCELoss(nn.Module): # Binary cross-entropy loss with masking def __init__(self): super().__init__() def forward(self, inputs, target, mask=None): out = nn.functional.binary_cross_entropy_with_logits( inputs, target, weight=mask, reduction="none") return out.mean(dim=1) loss = SigmoidBCELoss()

PyTorch 端之所以要自定义SigmoidBCELoss，是因为批量数据中含有大量填充位（padding），必须在损失计算时通过mask把它们排除在外。实现把 mask 作为weight传入binary_cross_entropy_with_logits，每个样本的损失为reduction="none"后按行求平均。

下面用给定变量验证损失的计算：第一行pred是模型的预测 logits，label标记正负例（第一行第 1 位为正例，第二行第 2 位为正例），mask表示哪些位置是有效词元（第二行只有前两位有效）。

#@tab all pred = d2l.tensor([[1.1, -2.2, 3.3, -4.4]] * 2) label = d2l.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]) mask = d2l.tensor([[1, 1, 1, 1], [1, 1, 0, 0]]) loss(pred, label, mask) * mask.shape[1] / mask.sum(axis=1)

第一行输出约为2.336，第二行约为1.654。上面的除法* mask.shape[1] / mask.sum(axis=1)把“按行平均”重新归一化为“在非掩码预测上平均”，相当于对每个样本去掉填充位的平均损失。

下面的手动计算（效率较低，但能看清原理）用 sigmoid 函数逐项验证了上述结果：两个输出可理解为在非掩码预测上平均的两个归一化损失。

#@tab all def sigmd(x): return -math.log(1 / (1 + math.exp(-x))) print(f'{(sigmd(1.1) + sigmd(2.2) + sigmd(-3.3) + sigmd(4.4)) / 4:.4f}') print(f'{(sigmd(-1.1) + sigmd(-2.2)) / 2:.4f}')

输出为2.3362与1.6540，与loss(...)的结果一致，验证了带掩码的二元交叉熵损失实现是正确的。

4.2 初始化模型参数

为词表中所有词分别定义两个嵌入层：一个用于词作为中心词的情形，一个用于词作为上下文词的情形。词向量维度embed_size设为 100：

#@tab mxnet embed_size = 100 net = nn.Sequential() net.add(nn.Embedding(input_dim=len(vocab), output_dim=embed_size), nn.Embedding(input_dim=len(vocab), output_dim=embed_size))

#@tab pytorch embed_size = 100 net = nn.Sequential(nn.Embedding(num_embeddings=len(vocab), embedding_dim=embed_size), nn.Embedding(num_embeddings=len(vocab), embedding_dim=embed_size))

两个嵌入层分别对应前向传播中的embed_v（net[0]）与embed_u（net[1]）。按词嵌入的惯例，训练完成后中心词向量net[0]用作最终词表示。

4.3 定义训练循环

由于数据中存在填充位，训练循环中的损失计算与一般分类任务略有不同：需要对pred按label的形状做 reshape，并在归一化损失时除以每个样本的有效词元数mask.sum(axis=1)。

#@tab mxnet def train(net, data_iter, lr, num_epochs, device=d2l.try_gpu()): net.initialize(ctx=device, force_reinit=True) trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': lr}) animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[1, num_epochs]) # Sum of normalized losses, no. of normalized losses metric = d2l.Accumulator(2) for epoch in range(num_epochs): timer, num_batches = d2l.Timer(), len(data_iter) for i, batch in enumerate(data_iter): center, context_negative, mask, label = [ data.as_in_ctx(device) for data in batch] with autograd.record(): pred = skip_gram(center, context_negative, net[0], net[1]) l = (loss(pred.reshape(label.shape), label, mask) * mask.shape[1] / mask.sum(axis=1)) l.backward() trainer.step(batch_size) metric.add(l.sum(), l.size) if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1: animator.add(epoch + (i + 1) / num_batches, (metric[0] / metric[1],)) print(f'loss {metric[0] / metric[1]:.3f}, ' f'{metric[1] / timer.stop():.1f} tokens/sec on {str(device)}')

#@tab pytorch def train(net, data_iter, lr, num_epochs, device=d2l.try_gpu()): def init_weights(module): if type(module) == nn.Embedding: nn.init.xavier_uniform_(module.weight) net.apply(init_weights) net = net.to(device) optimizer = torch.optim.Adam(net.parameters(), lr=lr) animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[1, num_epochs]) # Sum of normalized losses, no. of normalized losses metric = d2l.Accumulator(2) for epoch in range(num_epochs): timer, num_batches = d2l.Timer(), len(data_iter) for i, batch in enumerate(data_iter): optimizer.zero_grad() center, context_negative, mask, label = [ data.to(device) for data in batch] pred = skip_gram(center, context_negative, net[0], net[1]) l = (loss(pred.reshape(label.shape).float(), label.float(), mask) / mask.sum(axis=1) * mask.shape[1]) l.sum().backward() optimizer.step() metric.add(l.sum(), l.numel()) if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1: animator.add(epoch + (i + 1) / num_batches, (metric[0] / metric[1],)) print(f'loss {metric[0] / metric[1]:.3f}, ' f'{metric[1] / timer.stop():.1f} tokens/sec on {str(device)}')

几个实现要点：

• 优化器统一使用Adam，学习率lr=0.002；
• 设备默认通过d2l.try_gpu()自动选择 GPU（若无 GPU 则回退到 CPU），try_gpu定义于 d2l/torch.py；
• 训练指标d2l.Accumulator(2)累计“归一化损失之和”与“归一化损失个数”，最终损失为两者之比；
• 每轮训练中，animator在每 1/5 个 epoch 处记录一次平均损失，便于实时观察收敛。

在 PTB 数据集上执行训练（学习率 0.002，5 轮）：

#@tab all lr, num_epochs = 0.002, 5 train(net, data_iter, lr, num_epochs)

训练 5 个 epoch 后损失会明显下降，同时打印每个 epoch 的处理速度（tokens/sec）与运行设备。注意：这是教学规模的预训练，词向量质量与大规模语料上的完整 word2vec 仍有差距，但它验证了整条训练流水线是正确的。

5. 应用词嵌入：用余弦相似度检索语义相近词

预训练完成后，就可以用训练好的词向量做最经典的语义检索任务：找出与给定词余弦相似度最高的若干词。余弦相似度定义为向量夹角余弦：

$$\frac{\mathbf{x}^\top \mathbf{y}}{|\mathbf{x}| |\mathbf{y}|} \in [-1, 1].$$

由于词向量来自 skip-gram 的训练目标（中心词与其上下文词共现），语义相近的词往往具有相近的上下文分布，其向量夹角较小、余弦相似度较高。

#@tab mxnet def get_similar_tokens(query_token, k, embed): W = embed.weight.data() x = W[vocab[query_token]] # Compute the cosine similarity. Add 1e-9 for numerical stability cos = np.dot(W, x) / np.sqrt(np.sum(W * W, axis=1) * np.sum(x * x) + 1e-9) topk = npx.topk(cos, k=k+1, ret_typ='indices').asnumpy().astype('int32') for i in topk[1:]: # Remove the input words print(f'cosine sim={float(cos[i]):.3f}: {vocab.to_tokens(i)}') get_similar_tokens('chip', 3, net[0])

#@tab pytorch def get_similar_tokens(query_token, k, embed): W = embed.weight.data x = W[vocab[query_token]] # Compute the cosine similarity. Add 1e-9 for numerical stability cos = torch.mv(W, x) / torch.sqrt(torch.sum(W * W, dim=1) * torch.sum(x * x) + 1e-9) topk = torch.topk(cos, k=k+1)[1].cpu().numpy().astype('int32') for i in topk[1:]: # Remove the input words print(f'cosine sim={float(cos[i]):.3f}: {vocab.to_tokens(i)}') get_similar_tokens('chip', 3, net[0])

实现细节：

• 取中心词嵌入层权重W（形状为(词表大小, embed_size)），再取查询词chip对应的向量x；
• 余弦相似度通过批量点积除以范数乘积计算，分母加上1e-9防止数值不稳定导致的除零；
• topk(cos, k+1)取前 k+1 个最相似词，topk[1:]跳过第 0 位（即查询词本身，因为它与自己的相似度恒为 1）。vocab.to_tokens(i)把索引还原为词元字符串。

以'chip'为查询词、k=3 时，输出形如：

cosine sim=0.539: intel cosine sim=0.512: microprocessor cosine sim=0.491: chips

（具体数值会因训练随机性而略有差异。）可以看到，模型捕捉到了 “chip” 与半导体行业相关词汇（intel、microprocessor、chips）的语义关联——这正是词嵌入训练的目标体现。你也可以把查询词换成其他词（如'stock'、'bank'、'president'）来观察效果。

6. 核心要点小结

• skip-gram 模型可以只用嵌入层 + 批量矩阵乘法实现：中心词向量与上下文（含噪声）词向量做批量点积，得到形状(batch, 1, max_len)的相似度得分。
• 训练损失采用带掩码的二元交叉熵损失：masks排除填充位，labels区分正例（上下文词）与负例（噪声词）；归一化时要除以每个样本的有效词元数mask.sum(axis=1)。
• 模型参数是两个独立的嵌入层：一个用于中心词，一个用于上下文词；embed_size=100，优化器为 Adam（lr=0.002），训练 5 个 epoch。训练完成后，中心词向量net[0]作为最终词表示。
• word2vec 的典型应用是语义相似词检索：用余弦相似度对词向量排序，返回与查询词最相近的 k 个词（需剔除查询词自身）。

7. 延伸练习

1. 使用训练好的模型，对其他输入词查找语义相近词。试着调整超参数（如batch_size、max_window_size、num_noise_words、embed_size、lr、num_epochs），观察结果能否改善。
2. 当训练语料很大时，常采用“更新参数时再采样”的做法：即当前 minibatch 内，为每个中心词即时采样上下文词与噪声词，而不是一次性全部抽取。这样同一个中心词在不同训练轮次中会面对不同的上下文词或噪声词。想想这样做的好处（相当于更多的数据增广、更充分的负例多样性，同时降低内存占用），并尝试实现这种训练方式。具体来说，需要把数据集的中心词/上下文词/噪声词抽取逻辑从预处理阶段迁移到每次迭代内部。
3. 深入阅读相关章节可以继续：skip-gram 与 CBOW 的数学模型见 word2vec.md，数据预处理与 minibatch 装载细节见 word-embedding-dataset.md，负采样与分层 softmax 的理论推导见 approx-training.md，模型源码实现见 d2l/torch.py 与 d2l/mxnet.py。

赞分享

• 文档
• 教程
• 人工智能
• 深度学习
• NLP
• 计算机视觉
• 强化学习

【免费下载链接】d2l-en

Interactive deep learning book with multi-framework code, math, and discussions. Adopted at 500 universities from 70 countries including Stanford, MIT, Harvard, and Cambridge.

项目地址：https://gitcode.com/gh_mirrors/d2/d2l-en

点击查看免费下载

相关推荐

• d2l-en 词嵌入预训练数据集构建实战：从 PTB 原始语料到 skip-gram 负采样小批量

  d2l en 词嵌入预训练数据集构建实战：从 PTB 原始语料到 skip gram 负采样小批量 本文围绕《Dive into Deep Learning》

  文档教程人工智能深度学习NLP计算机视觉强化学习

• d2l-en 深度解析：word2vec 词嵌入原理、Skip-Gram 与 CBOW 模型及负采样训练实战

  d2l en 深度解析：word2vec 词嵌入原理、Skip Gram 与 CBOW 模型及负采样训练实战 导读 本文以《Dive into Deep Lea

  文档教程人工智能深度学习NLP计算机视觉强化学习

• 动手学深度学习（d2l-zh）实战：在 PTB 语料上使用负采样预训练 word2vec 跳元模型

  动手学深度学习（d2l zh）实战：在 PTB 语料上使用负采样预训练 word2vec 跳元模型 本文基于《动手学深度学习》中文版（d2l zh）仓库中的 c

  人工智能深度学习机器学习教程