简介基于PyTorch实现的词向量模型是一份面向自然语言处理初学者与深度学习入门者的代码资源。资源以Word2Vec的CBOW和Skip-gram思路为主线配套1803词的英文小型语料演示从文本分词、词汇表构建、One-Hot编码到MLP网络训练词向量的完整流程。通过PyTorch自动梯度机制可清晰观察损失下降与词向量空间变化适合用于教学实验或作为后续文本分类、情感分析、问答系统等任务的预处理基础。压缩包共6个文件包含3个py脚本、2个pyc缓存文件和1个txt数据集整体仅10KB结构紧凑便于快速下载与源码研读。其中Python代码覆盖数据读入、模型定义与训练调用等核心模块配合缓存文件可直接运行验证。目前已有3105人学习下载说明其在入门类资源中具备一定参考价值。1. 项目缘起与方案选型思考1.1 为什么用PyTorch做词向量模型词向量Word Embedding是自然语言处理里最基础、也是最绕不开的一环。2013年Word2Vec把词向量训练这件事带火之后各种实现层出不穷但真正上手自己训一个很多人反而会在这件事上卡住——不是不懂原理而是框架选型和工程细节没理顺。我选PyTorch的原因很朴素动态计算图让模型迭代验证非常舒服调试体验比静态图框架好太多了。你可以在训练过程中随时打印中间张量、修改模型结构不需要预先编译计算图。对比TensorFlow那一套尤其是早期版本PyTorch的Pythonic风格更贴近原生的研究直觉。虽然现在TensorFlow也支持了eager execution但生态偏向和社区活跃度上PyTorch在NLP领域明显占了主流。这套项目做的是一个基于PyTorch的标准词向量训练系统核心是复现Word2Vec的两种经典架构——CBOW和Skip-gram并在这个过程中把负采样、高频词下采样、学习率调度这些工程细节全部落地。不管是刚入门NLP的新手还是在研究里需要定制词向量的老手这套代码都可以直接作为基线使用。1.2 CBOW还是Skip-gram选型逻辑CBOWContinuous Bag-of-Words用上下文预测中心词Skip-gram反过来用中心词预测上下文。如果不考虑性能Skip-gram在大多数场景下学到的词向量质量要更好一些尤其是对低频词的表示更友好。原因在于Skip-gram在训练中会把中心词和每个上下文词都做一次配对预测等效于放大了低频词的训练次数每个样本的信息利用率更高。但代价是训练速度慢收敛需要更多epoch。CBOW则因为一次样本只做一次平均池化后预测计算量小训练速度快很多。项目里我把两个架构都封装了默认跑Skip-gram因为这套模型的应用场景是下游语义相似度计算对低频词的质量要求比较高。如果你只是做文本分类的特征提取换成CBOW会更划算。注意不要只盯着模型选型。词向量的最终质量受训练参数影响极其显著同样的语料下调整窗口大小和目标维度带来的差异往往比换架构更明显。2. 数据预处理与采样工程细节2.1 语料清洗和词表构建不管是哪个训练框架词向量的质量都严重依赖前置的数据流水线。原始语料里常见的噪声包括HTML标签、URL、非中英文字符、多余空白等。清洗这一步我分三层处理用正则表达式剔除无效字符统一全半角按标点和空白完成粗粒度分词中文场景建议配合jieba做细粒度分词统计词频过滤掉出现次数低于min_count的词词表构建这里有一个关键选型低频词直接删除还是用特殊标记替代。我建议直接把低频词扔掉。原因不是节省存储而是低频词的训练样本极少学出来的向量质量太差保留反而会给后续的语义计算引入大量噪声。min_count的默认值我设的是5语料质量高时调到3速度快很多。词表里还有一个细节需要在构建阶段就处理掉——为UNK未知词保留一个固定编号。训练完模型后新来的文本里必然有词表外的词如果不留UNK的位置推理时就只能报错。处理方式是在vocab里预留索引0给UNK索引1给PAD。2.2 负采样与高频词下采样的实现策略用原生softmax训练词向量有个致命问题输出层要计算词表大小的概率分布词表动辄几十万每次前向传播的矩阵乘法开销巨大。负采样的思路是把多分类问题转化成二分类问题——给定一组正样本和K组负样本让模型学会区分“真实出现的词对”和“随机抽出来的词对”。PyTorch里实现负采样不需要自己写复杂的采样逻辑一个F.logsigmoid就可以搞定。核心代码如下import torch import torch.nn as nn import torch.nn.functional as F class SkipGramModel(nn.Module): def __init__(self, vocab_size, embed_dim): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim) self.context_embedding nn.Embedding(vocab_size, embed_dim) # 初始化权重 nn.init.uniform_(self.embedding.weight, -0.5 / embed_dim, 0.5 / embed_dim) nn.init.uniform_(self.context_embedding.weight, -0.5 / embed_dim, 0.5 / embed_dim) def forward(self, center_word, pos_context, neg_context): # center_word: (batch_size,) # pos_context: (batch_size, num_pos) # neg_context: (batch_size, num_neg) center_emb self.embedding(center_word) # (batch, embed_dim) pos_emb self.context_embedding(pos_context) # (batch, num_pos, embed_dim) neg_emb self.context_embedding(neg_context) # (batch, num_neg, embed_dim) # 正样本得分 pos_score torch.sum(center_emb.unsqueeze(1) * pos_emb, dim-1) # (batch, num_pos) pos_loss -F.logsigmoid(pos_score).sum(dim1) # 负样本得分 neg_score torch.sum(center_emb.unsqueeze(1) * neg_emb, dim-1) # (batch, num_neg) neg_loss -F.logsigmoid(-neg_score).sum(dim1) return (pos_loss neg_loss).mean()高频词下采样Subsampling是另一个必要的预处理技巧。像“的”“了”“是”这类高频停用词它们在语料中出现的频率极高如果全部保留会在训练中占据大量样本导致其他词的表示空间被压缩。实践中按概率丢弃部分高频词概率公式为p(w) 1 - sqrt(t / f(w))其中f(w)是词的频率t是阈值通常设在1e-4到1e-5之间。f(w)远大于t的词被丢弃的概率就高f(w)接近t的词基本会被保留。3. 模型搭建与训练实操3.1 完整训练流程的工程实现数据管道处理完成后就进入模型训练环节。这里我不能只给一个孤零零的模型代码完整可复现的训练流程至少包括Batch采样器、负采样表构建、训练循环、模型保存。Batch采样器是最容易被忽视的环节。训练数据不能直接一次性加载到内存里——语料几百万行是常态。我采用的做法是按行流式读取语料动态构建训练样本。import random from collections import Counter import numpy as np import torch from torch.utils.data import Dataset, DataLoader class Word2VecDataset(Dataset): def __init__(self, corpus_path, window_size5, negative_num5, min_count5, sample_threshold1e-4): self.window_size window_size self.negative_num negative_num self.sample_threshold sample_threshold # 构建词表 word_counts Counter() for line in open(corpus_path, encodingutf-8): tokens line.strip().split() word_counts.update(tokens) word_counts {w: c for w, c in word_counts.items() if c min_count} self.vocab {w: i for i, (w, _) in enumerate(word_counts.items(), start2)} self.vocab[UNK] 0 self.vocab[PAD] 1 self.idx_to_word {i: w for w, i in self.vocab.items()} # 词频和采样概率 total_count sum(word_counts.values()) self.word_freq {self.vocab[w]: c / total_count for w, c in word_counts.items()} self.sample_probs {idx: 1 - np.sqrt(self.sample_threshold / freq) for idx, freq in self.word_freq.items() if freq self.sample_threshold} # 预计算负采样表 self.neg_table self._build_neg_table() # 预处理整个语料为句子列表小数据集 self.sentences [] for line in open(corpus_path, encodingutf-8): tokens line.strip().split() if len(tokens) 2: continue self.sentences.append([self.vocab.get(w, 0) for w in tokens]) def _build_neg_table(self): # Word2Vec论文中的负采样分布词频的3/4次幂 freq_table [] for word, idx in self.vocab.items(): freq self.word_freq.get(idx, 1e-5) freq_table.extend([idx] * int(freq ** 0.75 * 1000)) return freq_table def __len__(self): return sum(len(s) * self.window_size for s in self.sentences) def __getitem__(self, idx): # 通过索引定位到具体句子和词位置 for sent_idx, sentence in enumerate(self.sentences): sent_len len(sentence) - 1 if idx sent_len: idx - sent_len else: return self._generate_sample(sentence, idx) def _generate_sample(self, sentence, center_pos): center sentence[center_pos] context [] for offset in range(-self.window_size, self.window_size 1): if offset 0: continue pos center_pos offset if 0 pos len(sentence): context.append(sentence[pos]) if not context: return None context_idx random.choice(context) neg_idx random.sample(self.neg_table, self.negative_num) return torch.tensor(center), torch.tensor(context_idx), torch.tensor(neg_idx)负采样表用的是原始论文中的3/4次幂加权分布把词频的0.75次方作为采样权重。这个幂指数能让低频词被抽中的概率相对提高从而改善低频词的训练效果。代码里把这个表构造成了一个大的list每次随机从中抽出负样本。词表大的时候这个list可能会占几百MB内存需要注意。3.2 关键技术参数的选择依据训练词向量的参数配置直接决定输出质量。我整理了一份项目里测试过的参数组合可以直接参考使用参数推荐值说明embedding维度100~300维度越高表达能力越强但需要更多语料否则过拟合窗口大小5语义相近任务用5语法相近任务可到10负样本数5~10语料大时可以用5小的语料建议10学习率0.025Word2Vec原论文设置后期做线性衰减最小词频5低于该频次的词直接剔除batch size256~1024显存够大时尽量大训练稳定epoch数5~10早期epoch不敏感后期需配合学习率衰减学习率的设置值得多说两句。Word2Vec的常见陷阱是词向量训练不像分类任务那样有明确的loss收敛标准过大的学习率会让词向量在空间里分布过于离散过小则会欠拟合。我在项目中采用的方案是“初始0.025 指数衰减”每个epoch结束后乘以0.8到第5个epoch基本稳定在0.008左右。实测下来这个曲线能兼顾低层语义空间的展开和高层微调。模型训练到这里还不能直接结束——embedding层和context embedding层的选择很关键。从输入侧学到的embeddingself.embedding才是最终要保存的词向量context embedding只是训练过程中的副产品。要在保存模型时明确这一点避免导出错了张量。实操心得PyTorch里Embedding层默认初始化是标准正态分布但在词向量训练场景下建议手动改成均匀分布初始化和更小的初始范围。原因在于词向量的最终分布本质上是个低维流形初始范围太大会让模型在早期需要大量迭代才能把向量“拉回”合理的语义空间延长收敛时间。4. 常见问题与排查技巧实录4.1 loss不下降或收敛过慢词向量模型最常见的困境是跑了几个epochloss曲线纹丝不动或者下降幅度远低于预期。多数时候问题不在模型结构上而在数据采样环节。首先要检查负采样表是否构建正确。如果负采样表只包含少数高频词模型很快就会陷入“区分不了正负样本”的困境loss直接卡住。我在代码里构建负采样表时曾经把词频取幂的符号写反了——频率用负幂次导致低频词占了绝大多数负样本模型学到第3个epoch就崩了。排查办法很朴素打印负采样表的前20个词看看是不是符合“含少量常见词大量中频词”的分布特征。其次是数据管道是否有空样本进入训练。__getitem__里如果返回了None而DataLoader没有过滤会导致batch形状不对训练直接报错。项目里写了一个collate_fn来过滤空样本这个函数在数据管道中是必选项而不是可选项。4.2 词向量余弦相似度结果异常模型训练完成表面看起来loss很低但实际检索“北京”的近义词返回的全是“中国”“首都”这类关系词而不是“上海”“广州”这类实体对等词。这个现象很多人会误解为模型训练失败其实是评估方式误导了自己。Word2Vec学到的是分布相似性不是实体齐性。“北京”和“中国”经常出现在同一上下文中“北京是中国的首都”两者的向量方向确实很接近。如果想得到“同类城市”的向量更合理的做法是直接在原始embedding上做word analogy类比推理比如“北京-中国日本”接近“东京”。这个评测方式在gensim里能直接用好过拍脑袋看相似度。4.3 PyTorch环境安装与版本匹配问题我在这套项目上迭代过几个版本环境配置踩过的坑比模型本身的还多。主要问题集中在PyTorch版本和CUDA版本的匹配上。这里给出一个很实用的建议不要装最新版。PyTorch官方每半年发一个大版本但很多第三方库比如transformers、mmcv系列的兼容节奏要慢半拍装完新版本后跑旧代码经常报出莫名其妙的cuda错误。我现在的组合是Python 3.10 PyTorch 2.1.x CUDA 11.8覆盖绝大多数应用场景稳定运行了近一年。安装时注意区分CPU版和GPU版。CPU版做词向量训练完全能跑只是速度慢一些如果只是学习原理可以先从CPU版开始。安装命令是# CPU版 pip install torch torchvision torchaudio # GPU版CUDA 11.8 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu1184.4 显存溢出与训练效率问题词向量模型本身不复杂但词表超过50万时负采样列表会膨胀到数GB加上embedding矩阵、训练数据内存开销不容小觑。当训练过程中出现内存/显存溢出先从这几个方向排查负采样表构建方式的优化空间。不要用list存储改成numpy.array的int32类型内存能省一半batch size是否设置过大。虽然大batch能加速但词向量训练其实不需要极大batch256就够数据加载是否用了多进程。DataLoader(num_workers4)能显著加速数据读取但Windows系统下num_workers要设为0否则会死锁5. 模型评估与应用扩展5.1 定性定量结合评估词向量质量训练完成后不要急着保存下线先做一轮完整的质量评估。最常用的方式是“相似度任务类比任务”的组合评估。相似度任务构建一个小的“词语对人工评分”集合计算模型cosine相似度与人工评分的相关性。传统上用的是WordSim353数据集中文场景可以用自己标注的几百条词对做快速验证。相关性指标用Spearman相关系数。类比任务直接复用word analogy的形式。代码实现如下def analogy(model, word, relation_word, target, topn10): # word - relation_word target word_emb model[word] relation_emb model[relation_word] target_emb model[target] query_emb word_emb - relation_emb target_emb # 与词表中所有词计算相似度 all_embeddings torch.tensor(model.vectors) scores torch.matmul(all_embeddings, query_emb) / ( torch.norm(all_embeddings, dim1) * torch.norm(query_emb) ) top_indices torch.topk(scores, topn).indices return [model.index_to_key[i] for i in top_indices]评估结果是调整训练参数的直接依据。如果相似度任务得分低优先考虑增大语料量和调整窗口大小如果类比任务得分低优先考虑增大向量维度或负样本数。5.2 从词向量到下游任务平均池化与模型融合训好词向量之后下一步通常是把词向量应用到具体NLP任务。最简单的用法是把句子中的词向量取平均作为整个句子的表示然后喂到分类器里。def sentence_to_vector(sentence, model, tokenizer, embedding_dim100): vectors [] for token in tokenizer(sentence): if token in model: vectors.append(model[token]) if not vectors: return np.zeros(embedding_dim) return np.mean(vectors, axis0)这个平均池化方法简单但非常有效适合作为文本分类、相似度计算的基线。如果你的任务中还有上下文信息可以利用可以在词向量之上再接BiLSTM或Transformer层。不过这里要提醒一句预训练词向量在较大模型里只适合作为embedding层的初始化不要冻结整个embedding层。冻结会导致模型无法根据任务微调词向量反而限制了效果。实践中设定embedding.weight.requires_grad True让模型在反向传播中微调词向量往往能带来2~3个百分点的效果提升。模型融合是另一个可以操作的扩展点把CBOW和Skip-gram分别训练出的词向量拼接或者将多个epoch的检查点进行加权融合能提升词向量的稳定性和泛化能力。原因是两个模型的优化路径不同学到的是语义空间的不同视角拼接后等效于增大了表示容量。6. 个人实操体会词向量这个项目看似简单真要把效果做到可用级别需要打磨的细节远超预期。我在实际迭代中最深的感受是语料清洗和采样策略对结果的影响比模型结构的调整大得多。很多人花大量心思改网络结构却忽视了语料的覆盖度和质量导致最后的语义空间扭曲。另外PyTorch在这里提供了足够的灵活度去实现自定义逻辑。常规的深度学习教程把重心放在网络层搭建上但词向量训练的工程核心其实在数据管道负采样表的构建、采样概率的计算、动态读取与流式训练这些环节的代码量比模型本身多得多。如果读者在这几个模块上多花时间调试后续做任何NLP项目都会顺手很多。项目后续可以考虑的扩展方向是接入大规模语料做增量训练、支持多语言混合语料、引入动态词向量如ELMo的训练思路。这些都是在当前代码基线之上做的自然延伸跑通的主体框架不需要动。本文还有配套的精品资源点击获取