简介一份面向NLP学习者的PyTorch词向量模型实现资源覆盖Word2Vec的CBOW与Skip-gram核心思路适合希望从零理解词嵌入原理并动手实践的开发者和学生。资源共6个文件包含3个Python源码脚本、2个编译后的pyc缓存文件及1个英文文本数据集内含1803个单词压缩包仅10KB轻量易得。当前已有3110人浏览学习。内容完整演示了从文本分词、词汇表构建、One-Hot编码到神经网络训练、损失计算与权重保存的全流程并展示了PyTorch自动求导机制如何简化模型训练。配套的小型语料可用于教学实验训练所得词向量能应用于文本分类、情感分析和检索等下游任务。代码结构清晰适合作为NLP入门实践或深度学习框架学习的参考样例。1. 基于PyTorch重建词向量模型不调现成库自己训练一个word2vec之前我接手过一个领域文本分类的需求语料里全是产品型号和行业黑话套用通用词向量模型怎么调都泛化不上去。后来索性自己用 PyTorch 搭了一个词向量训练管线自己洗数据、自己建词表、自己用负采样训练 skip-gram一晚上就能把整个项目的 embedding 重新刷一遍。这个仓库就是那套管线的整理版本核心是三条链路——文本清洗与词表构建、skip-gram 负采样模型、训练与检索验证跑通之后就能得到一份领域专属的词向量表。它适合想把词向量用在自己数据上、又不想引一堆重型依赖的自然语言处理入门者也适合需要在特征层做可控实验的机器学习开发者。2. 项目架构与数据流从原始文本到 mini-batch 的完整链路2.1 目录结构与模块职责拿到仓库之后我习惯先看目录对工程结构有个底。这套实现的目录划分得非常直白读代码的路径是dataset.py - model.py - train.py - evaluate.py每一步都对应训练流程中的一个环节。word2vec-pytorch/ ├── data/ │ ├── text8/ # 样例英文语料 │ └── vocab.json # 训练完成后导出的词表 ├── src/ │ ├── dataset.py # 语料读取、词表构建、负采样表、mini-batch 生成 │ ├── model.py # SkipGram / CBOW 模型定义 │ ├── train.py # 训练入口负责配置加载与循环 │ └── evaluate.py # top-k 相似词、类比推理验证脚本 ├── config.py # 全部超参数集中管理 ├── requirements.txt └── README.mdconfig.py把 embedding 维度、窗口大小、负采样数量、epoch、min_count、batch_size、学习率全部放在一起。我自己的习惯一直是“配置和代码分离”这种小项目可能觉得无所谓但改超参数做对比实验时集中管理会比在代码里到处翻要省事得多。2.2 数据流从原始文本到 mini-batch 的四步转换这套项目把数据流转拆成了四个明确的阶段每一步的产物都会在内存里被下一个模块消费。第一步是语料读入与分词。英文语料直接按空白字符 split 就能用中文语料则要先经过分词器处理成“空格分隔的 token 序列”再进入同一套管线。dataset.py统一接收List[List[str]]这种形态也就是说你只要能把文本变成词的列表后面的逻辑就不用动。第二步是词表构建。全量统计 token 频次之后会把频次低于min_count的词直接丢弃。这一步过滤掉的通常是拼写错误、太稀有的命名实体保留下来的是有统计意义的词词表大小也会从几十万缩到几万训练速度会明显变快。第三步是高频词亚采样subsampling。这一步是用一个概率把高频停用词从训练样本中踢掉公式是p(w) 1 - sqrt(t / f(w))t默认取1e-4。我一开始没开这个选项训练出来的向量里“the”“of”这类词几乎占据了相似词列表的前十名开了之后情况立刻改善。第四步是滑动窗口生成训练对。对每个中心词在窗口范围内取上下文词组成(center, context)正样本对。窗口大小不是固定的每次会从 1 到window之间随机抽一个值这是原论文里的做法相当于给远距离词对降权。2.3 关键技术决策独立双 Embedding 与负采样项目里有一个值得注意的设计中心词和上下文词各用一张独立的 Embedding 表。训练过程中同一个词既会被放到 center 位置也会被放到 context 位置如果两张表共用参数正负梯度会在同一份权重上互相干扰收敛会变慢。分表之后训练结束只取center_embed.weight作为最终词向量context 表只承担训练阶段的辅助作用。损失函数选择的是负采样negative sampling而不是原生 softmax这一点对工程实现影响很大。词表动辄几万每次训练都要算一次全词表的归一化概率计算量完全吃不消负采样把多分类问题拆成了“给定中心词预测上下文词”和“给定中心词排除噪声词”两类二分类问题每个 batch 只需要计算 K 个负样本的得分训练成本直线下降。项目里默认 K5与论文推荐的取值范围一致。3. 核心代码逐段拆解词表构建、负采样表、模型与训练循环3.1 词汇表构建min_count 过滤与 id 映射词表是整条训练链路的地基dataset.py里的Vocab类承担了词频统计和 id 映射两件事。from collections import Counter class Vocab: def __init__(self, tokenized_corpus, min_count5): counter Counter() for tokens in tokenized_corpus: counter.update(tokens) self.word2id {} self.id2word {} self.word2freq {} idx 0 for word, freq in counter.most_common(): if freq min_count: continue self.word2id[word] idx self.id2word[idx] word self.word2freq[word] freq idx 1 self.vocab_size idx这里counter.most_common()返回的是按频次降序排列的列表遇到第一个低于min_count的词就可以 continue 而不是 break因为后面只会更低所以写 continue 逻辑上不会出错只是多扫几轮。id 分配顺序与频次相关高频词排在前面低频词排在后面这种排布在后续生成负采样表时遍历顺序更自然。word2freq字典必须保留它是负采样和亚采样的直接依赖。如果在构建词表时把频次信息丢掉后面就要回头重新统计一遍内存和时间都亏。3.2 负采样表基于词频的 0.75 次幂驯化负采样不是均匀随机抽词而是按词频的 0.75 次幂加权抽取。0.75 次幂是论文里验证过的经验值作用是把“的”“了”这类超高频词的采样概率拉低给中频词更多被当成负样本的机会训练出的向量对低频词更友好。import random def build_neg_table(vocab, power0.75): table [] total sum([freq ** power for freq in vocab.word2freq.values()]) table_size int(1e7) for word, freq in vocab.word2freq.items(): count int((freq ** power) / total * table_size) table.extend([word] * count) return table def sample_neg(neg_table, k, exclude_set): candidates [] while len(candidates) k: word random.choice(neg_table) if word not in exclude_set: candidates.append(word) return candidates这里先把词频做 0.75 次幂归一化再按比例映射到总长度 1e7 的采样表里。用table.extend([word] * count)这种方式把每个词的 id 复制多份采样时random.choice就是 O(1) 操作。内存上每个 int 约 8 字节1e7 个整数大约 80MB在可接受范围。注意sample_neg里的exclude_set它应该包含当前样本的上下文词集合。如果负样本里混进了真正的上下文词模型会学着把正样本也判成负样本loss 就会出现诡异的波动。这是我踩过的坑后面排查章节会展开说。3.3 模型定义双 Embedding 结构与矩阵乘法打分model.py里的SkipGramModel实现非常紧凑forward 里只出现两类操作查 Embedding 表然后做 batched matrix multiplication。import torch import torch.nn as nn class SkipGramModel(nn.Module): def __init__(self, vocab_size, embed_dim128): super().__init__() # 中心词和上下文词各用一张独立的嵌入表 self.center_embed nn.Embedding(vocab_size, embed_dim) self.context_embed nn.Embedding(vocab_size, embed_dim) init_range 0.5 / embed_dim self.center_embed.weight.data.uniform_(-init_range, init_range) self.context_embed.weight.data.uniform_(-init_range, init_range) def forward(self, center, context, negatives): center_vec self.center_embed(center) # (B, D) center_vec center_vec.unsqueeze(1) # (B, 1, D) context_vecs self.context_embed(context) # (B, C, D) neg_vecs self.context_embed(negatives) # (B, K, D) pos_affinity torch.bmm( context_vecs, center_vec.transpose(1, 2) ).squeeze(-1) # (B, C) neg_affinity torch.bmm( neg_vecs, center_vec.transpose(1, 2) ).squeeze(-1) # (B, K) return pos_affinity, neg_affinitycenter_vec.unsqueeze(1)是把中心词向量从(B, D)变成(B, 1, D)这样torch.bmm才能把每个样本的上下文词向量矩阵和中心词向量做批量内积。bmm(context_vecs, center_vec.transpose(1, 2))实际算的是(B, C, D) x (B, D, 1)输出(B, C, 1)最后的squeeze(-1)去掉最后一个维度。初始化用的是uniform_(-0.5/embed_dim, 0.5/embed_dim)数值范围很小。一开始我调大过初始化范围发现第一个 batch 的 loss 会冲到 8 以上收敛速度明显变慢。Embedding 初始化和线性层的初始化逻辑不一样不能用默认的均匀分布直接套。3.4 训练循环采样、前向、损失与参数更新训练循环本身不长但放了很多细节。负采样 loss 使用的是log-sigmoid正样本打分希望趋近 1负样本打分希望趋近 0。import torch.nn.functional as F def train_one_batch(model, center, context, negatives, optimizer): pos_score, neg_score model(center, context, negatives) pos_loss -F.logsigmoid(pos_score).mean() neg_loss -F.logsigmoid(-neg_score).mean() loss pos_loss neg_loss optimizer.zero_grad() loss.backward() optimizer.step() return loss.item()-F.logsigmoid(pos_score)对应的是“正样本被判定为 1 的负对数概率”-F.logsigmoid(-neg_score)对应的是“负样本被判定为 0 的负对数概率”。两者直接相加就是这一个 batch 的总损失不需要额外加权因为负样本的数量已经通过采样次数控制了。正样本对是滑动窗口产生的每一步向前取 C 个上下文词所有上下文词赋予相同的监督信号。一个词的语义是从它的“共现伙伴”身上学到的窗口大小直接决定了共现半径窗口大学到的更多是主题邻近关系窗口小学到的更多是语法搭配关系。项目默认窗口为 5这也是语言建模场景里比较折中的一个取值。4. 训练与验收loss 曲线之外如何验证 embedding 真的学到了语义4.1 训练参数说明与收敛判断config.py里的默认超参数组合如下参数默认值说明embed_dim128向量维度增大则语义容量提升但训练变慢window5上下文窗口半径实际取 1~5 随机值negative5每个正样本对对应的负样本数min_count5词频低于该值的词直接丢弃subsample_t1e-4高频词亚采样阈值越小过滤越狠batch_size256每个 mini-batch 的样本数量lr0.025初始学习率训练日志里主要盯 loss 的走势。负采样模式下正常情况 loss 会从 0.7 左右快速下降前几个 epoch 能压到 0.3 上下之后进入缓慢下降阶段。loss 并不追求降到特别低如果压到 0.05 以下反而要警惕是不是模型把所有负样本都错误地当成了正样本。比较稳的做法是把每个 epoch 的平均 loss 打到日志文件里同时每个 epoch 结束后跑一遍 evaluate 脚本人工看一眼相似词质量。loss 只告诉你模型在拟合训练目标不告诉你 embedding 的语义质量。4.2 相似词与类比推理验证evaluate.py里实现了一个朴素的 top-k 相似词查询原理是取目标词的向量和词表所有向量做余弦相似度取最高的一批 id。def most_similar(words, vocab, model, topk10): ids [vocab.word2id[w] for w in words if w in vocab.word2id] vecs model.center_embed.weight.data[ids].mean(dim0) all_vecs model.center_embed.weight.data sims all_vecs vecs / (all_vecs.norm(dim1) * vecs.norm() 1e-9) topk_ids sims.topk(topk 1).indices.tolist() return [(vocab.id2word[i], sims[i].item()) for i in topk_ids]mean(dim0)是为了支持多词类比输入比如传入[king, woman, man]先把三个词的向量取平均再和全表做相似度计算排名靠前的词大概率就是类比目标。这种“加减向量找最近邻”的验证方式比单看词本身更能反映向量空间是否学到了语义结构。跑相似词验证时我会拿同一批词在不同的 epoch 轮次做对比。比如训练 5 个 epoch 后查“computer”的邻居出来的还是“pc、laptop”这类强共现词训练 15 个 epoch 后语义上更远但主题相关的词也会挤进前十。这是一个非常直观的质量提升信号。4.3 t-SNE 可视化和异常词检查除了数值指标可视化是发现 embedding 问题的利器。做法很简单从词表里挑 500~1000 个出现频次较高的词把它们的向量用 t-SNE 降到二维然后散点图上看聚类结构。理论上语义相近的词应该聚成团比如体育类、金融类、技术类各占一片区域。最容易观察到的异常是大量停用词聚成一团悬浮在图的中央这通常意味着亚采样阈值没生效或者 min_count 设置太低。另一个异常是有明显语义区别的词混在同一个簇里这多半是语料本身有噪声——比如爬虫文本里混了大量 HTML 标签分词之后这些符号词成了高频 token就会拉偏向量空间。我自己的习惯是训练收尾前跑一次可视化觉得簇边界清晰了再导出向量文件。可视化不是必须步骤但它能在几秒钟内暴露很多数值指标看不到的问题。5. 跑项目时的常见翻车现场与排查记录5.1 现象一训练到一半 loss 突然变成 0有次训练端到端跑了三个小时结果从日志上看 loss 从第 50 个 batch 开始直接变成 0。相似词查询结果也全乱了所有词的 top 邻居都是同一个词。当时第一反应是学习率太大导致梯度爆炸调小了也没用。排查后发现原因是负采样函数里exclude_set写错了。那个集合在样本生成时只包含了当前中心词本身没有包含窗口内的上下文词于是上下文词也被当成负样本抽了出来。模型收到“同一个词既该预测为正又该预测为负”的矛盾信号梯度互相抵消loss 就被压到了 0。解决方式是把sample_neg的排除集合从{center}改成set(context_tokens [center])确保负样本不会包含窗口内的任何词。从那以后我每次检查数据生成代码都会先确认正负样本的边界是否真的互斥。5.2 现象二加载训练好的向量维度对不上训练完向量后下游模型加载时报RuntimeError: Expected embedding weight size [50001, 128], got [49998, 128]。对比词表后发现下游任务的测试集里多了三个训练语料里没出现过的词。这是自然语言处理项目里最容易犯的错词表的构建必须统一训练和预测各自建词表必然出现边界词问题。解决方式是训练的时候把vocab.json和向量文件一起导出下游加载时先读词表再按词表顺序读向量。遇到新词常见做法是用weight.new_empty(embed_dim).uniform_(-0.5/embed_dim, 0.5/embed_dim)初始化一个随机向量塞进去。虽然随机初始化的新词向量没有语义信息但至少保证了训练和推理的计算图能跑通。5.3 现象三相似词全是清一色的高频停用词训练完查“人工智能”的相似词出来的结果是“的”“是”“在”——这种结果基本等于白训。原因有两层第一是min_count设太低大量只出现一两次的噪声词参与了训练把低频词的向量空间搅乱了第二是没开亚采样停用词在滑动窗口里出现的频率远高于内容词在大量上下文里反复充当上下文词自然学会了和所有词都“搭边”。解决方式是先调min_count从 5 开始向上试探同时开启 subsampling。效果立竿见影停用词的干扰会大幅下降。如果调完还出现低频词相似度普遍偏低的情况可以把负样本数从 5 提到 10让低频词在训练中接收更严格的反向信号。5.4 现象四显存或内存占用过大小语料训练一般不会碰到这个问题但一旦语料换成几 GB 级别的文本负采样表和 batch 张量就会吃掉大量内存。负采样表 1e7 长度占 80MB如果再把power调成 1 或者把table_size调大内存会翻倍。另一个常见的资源黑洞在torch.bmm上。context_vecs张量的形状是(B, C, D)其中C是窗口内上下文词数量。如果把 batch_size 拉到 1024window 设为 5负样本设为 15一个 batch 里context和negatives两个矩阵的中间结果会占用相当可观的显存。解决方式是调低 batch_size 到 128 或 256或者把负样本数从 15 降到 5训练速度并不会明显变差但显存能省下一大半。6. 二次开发把词向量模型改成你自己的 NLP 前置模块6.1 换语料中文分词与字符级向量的接入点这个项目的原始代码默认处理的是按空格切好的英文 token 序列。接入中文语料时不需要动模型结构只需要在数据入口处把分词结果转换成空格分隔的文本。常见的做法是用现成的分词库在dataset.py里增加一个 tokenizer 参数def load_corpus(raw_path, tokenizerNone): lines [] with open(raw_path, encodingutf-8) as f: for line in f: line line.strip() if tokenizer: tokens tokenizer(line) else: tokens line.split() if tokens: lines.append(tokens) return lines如果做的是电商场景还可以在分词之后顺便做品牌名合并、型号归一化。这一步放在词表构建之前会让最终向量更贴合业务术语。6.2 改结构从 Skip-gram 切换到 CBOW项目里同时保留了两种模型的接口改起来非常直接。CBOW 和 Skip-gram 的差别只在输入输出方向上CBOW 用上下文词预测中心词forward 里把context_vecs在维度 1 上取均值再和中心词向量做内积打分。context_mean context_vecs.mean(dim1) # (B, D) score (context_mean * center_vec).sum(dim-1)切换之后只需要把数据流从“中心词 - 上下文词”改成“上下文词 - 中心词”其余训练基础设施完全可以复用。两个方向对比着训会发现在同一份语料上Skip-gram 对低频词更友好CBOW 的训练速度更快选哪个取决于下游任务的语料分布。6.3 往外走把 embedding 接进下游任务词向量训练的终点不是得到一个 .txt 文件而是作为下游任务的初始化参数。做法是把center_embed.weight转成nn.Embedding的初始化权重然后选择冻结或者微调。领域数据充足的场景建议微调冷启动场景建议先冻结让上层任务先收敛再解冻微调。从那以后我每次做自然语言处理项目都会先用自己的这套脚本把领域语料过一遍。训练一次只需要一顿饭的时间却能让下游模型的特征输入精准命中业务词汇的分布。希望这个项目解析能帮你把词向量这一步跑顺手少走我踩过的那些弯路。本文还有配套的精品资源点击获取