简介一份基于PyTorch实现的词向量模型代码包面向自然语言处理初学者以及希望动手理解Word Embedding的开发者。内容围绕CBOW、Skip-gram等经典思路展开演示如何把词汇映射为连续向量、构建词表并搭建简单网络完成训练可作为文本分类、情感分析等下游任务的预处理练习起点。压缩包共6个文件以3个py源码为主另有2个pyc编译缓存与1个txt英文语料源码负责数据处理与模型定义txt语料则提供含1803个单词的小型训练集整体仅约10KB便于本地快速跑通。已有3111人学习下载。篇幅虽小却覆盖数据预处理、模型构建、交叉熵损失计算与权重保存的完整链路并保留训练后的应用示例适合边读代码边调向量维度、窗口大小与迭代次数直观观察词向量质量变化再迁移到其他自然语言处理任务中。1. 从一份纯文本到一个能算距离的坐标PyTorch 词向量模型交付什么拿 PyTorch 手写一份词向量模型难的从来不是公式是让它在自己的语料上跑出像样的相似词。有个做评论情感分类的开发者标注数据只有三千条分类器怎么调都在 78% 徘徊后来他用几十万条无标注评论先训了一份词向量把句子里的词换成 100 维稠密向量再喂给分类器同一份标注数据准确率往上走了六个点。词向量模型干的不是业务是把「词」这种离散符号压成坐标让机器学习模型能算距离、做加减法。这份基于 PyTorch 实现的词向量模型走的是自然语言处理里最经典那条路纯文本输入Skip-gram 和 CBOW 两套结构带负采样与层次 Softmax训完能查相似词、做类比推理、导出可视化。想亲手把 word2vec 从公式写成代码的或者要拿自己领域语料重训一份向量的都合适词表从几万到几十万单卡甚至纯 CPU 都能跑。2. 项目骨架与技术选型七个 Python 文件撑起一次完整训练2.1 目录结构与各模块职责划分拆这类训练项目我第一眼看的是目录有没有把「数据加工」和「模型计算」分开。分不开的项目改一次分词规则就得动训练脚本后面基本没法维护。这份实现的分层是干净的word2vec-pytorch/ ├── configs/ │ └── sg_zh.yaml # 一次实验一份配置超参不写死在代码里 ├── data/ │ ├── raw/corpus.txt # 原始语料一行一段不做任何预处理 │ └── processed/ │ ├── tokens.txt # 分词结果词之间空格分隔 │ └── vocab.json # 词表与词频训练和推理共用同一份 ├── w2v/ │ ├── vocab.py # 词表构建、ID 映射、低频词截断 │ ├── dataset.py # 滑窗采样产出 (center, context) 批次 │ ├── model.py # SkipGram / CBOW 两个 nn.Module │ ├── trainer.py # 训练循环、日志、checkpoint 落盘 │ ├── evaluate.py # 相似词、类比推理、t-SNE 可视化 │ └── utils.py # 负采样表、随机种子、向量归一化 ├── scripts/ │ ├── preprocess.py # raw/corpus.txt - tokens.txt vocab.json │ ├── train.py # 入口读 yaml 开跑 │ └── export.py # 导出 .vec / .npy 给下游任务 ├── requirements.txt └── README.mdvocab.py是唯一持有「词 → ID」映射的地方训练和推理都从它拿避免了两边词表对不上的血泪事故。dataset.py只负责产出整数 ID不碰字符串这样后面换分词器不用改采样逻辑。model.py里两个 Module 都很薄核心就两张 Embedding 表加一个前向损失把复杂度压在损失函数而不是网络结构上这是 word2vec 这类模型的特点——网络浅但采样和损失讲究。trainer.py单独抽出来是为了能在不碰模型定义的前提下替换优化器、加梯度累积、加日志钩子。evaluate.py属于「训完必须马上能验证」的那部分很多实现把它漏掉结果训练跑完几小时想查个相似词还得现写脚本。configs/*.yaml这层看着多余实际是省后悔药的地方。词向量训练对一个参数极其敏感——窗口大小5 和 10 训出来的向量在类比任务上能差一截。把window、neg_k、min_count、subsample全部外置到配置文件一组实验一个文件比在代码里注释来注释去靠谱得多。2.2 依赖环境与最小可运行配置依赖不多锁死版本能省掉一堆玄学问题torch1.13,2.4 numpy1.21 tqdm pyyaml jieba # 中文语料才需要纯英文语料可以删掉 scikit-learn # 只给 t-SNE 可视化用不装也能训练Python 用 3.8 到 3.11 都行3.12 上部分老版本 torch 轮子会缺装之前先确认。CPU 机器直接pip install torch --index-url https://download.pytorch.org/whl/cpu别装默认包默认包会拖一个几百兆的 CUDA 运行时下来纯浪费。有卡的按 CUDA 版本对号入座装错了报的不是「版本不匹配」而是CUDA error: no kernel image is available for execution on the device这行报错能让人查半天。scikit-learn只在evaluate.py --plot时才 import写成函数内延迟导入这样在只有 torch 的环境里训练和查相似词照样能跑。jieba同理英文语料不需要它缺了也别让整个包 import 失败。提示先把data/raw/corpus.txt换成一份几千行的小语料跑通全流程再换大语料。大语料上出的问题和小语料上出的问题排查成本差十倍。2.3 一次训练从读取文本到落盘向量的完整链路整条链路是五步每一步的产物都能单独检查这点很重要scripts/preprocess.py读data/raw/corpus.txt做归一化和分词写出tokens.txt和vocab.json。w2v/vocab.py加载vocab.json把tokens.txt每行映射成 ID 序列。w2v/dataset.py在 ID 序列上滑窗攒够batch_size个(center, context)对就吐一批。w2v/trainer.py每个 batch 额外采neg_k个负样本前向、反向、更新按 epoch 存 checkpoint。scripts/export.py只导出in_emb输入向量表成.npy或文本格式供下游nn.Embedding.from_pretrained用。第 2 步和第 5 步是最容易出岔子的地方第 2 步决定词表顺序第 5 步决定下游能不能对上号。所以vocab.json里除了词表本身还要存一份id2token的数组顺序就是 Embedding 的行号导出时原样带出去。下游加载向量矩阵时行号必须和这份数组一一对应对不上就是「查出来的相似词毫无关联」这类问题且极难定位。链路中间产物tokens.txt建议保留别训完就删。语料里混进乱码、分词把「神经网络」切成「神经」「网络」这些问题都是回看这个文件才发现的。3. 语料预处理与词表构建从原始文本到 (center, context) 样本对3.1 分词、清洗与低频词截断预处理这步决定上限。词向量对语料质量极其敏感喂进去一堆 URL、表情符号、重复刷屏文本训出来的向量必然脏。import re import jieba def normalize(line: str) - str: line line.strip().lower() # 数字统一成占位符避免 num 之外的千奇百怪形态污染词表 line re.sub(r\d, num , line) # 只保留中英文、数字占位符和基本词字符其余一律替换成空格 line re.sub(r[^\w\u4e00-\u9fff], , line) return re.sub(r\s, , line).strip() def tokenize(line: str): # 英文按空格切中文交给 jieba实际项目里英文还要过一遍停用词 tokens [] for seg in line.split( ): if not seg: continue if re.search(r[\u4e00-\u9fff], seg): tokens.extend(w for w in jieba.lcut(seg) if w.strip()) else: tokens.append(seg) return tokensnormalize里lower()对中文没影响对英文是必须的否则Model和model会变成两个词各训一份向量两边都学不充分。数字替换成num是个取舍不替换的话年份、价格、ID 会撑爆词表且每个出现次数都很少替换之后这些小样本被归并成一个高频 token反而让上下文窗口更干净。tokenize里判断一段是不是中文再决定用不用 jieba是为了处理中英混排。纯英文语料整段走空格切分就够了比先 jieba 再拼回去快得多。停用词这块我的习惯是不在预处理阶段删而是在词表构建时通过min_count自然过滤——「的」「是」这类词频高删不掉但它们在负采样里会被降权后面会讲比硬删更符合 word2vec 的设计。低频词截断在vocab.py里做from collections import Counter class Vocab: def __init__(self, min_count5, max_sizeNone): self.min_count min_count # 词频低于此值的词直接丢弃 self.max_size max_size # 词表上限超出按词频截断 self.token2id {} self.id2token [] self.counts [] def build(self, token_lines): counter Counter() for tokens in token_lines: counter.update(tokens) items [(w, c) for w, c in counter.items() if c self.min_count] items.sort(keylambda x: -x[1]) # 按词频降序保证可复现 if self.max_size: items items[:self.max_size] self.id2token [w for w, _ in items] self.token2id {w: i for i, w in enumerate(self.id2token)} self.counts [c for _, c in items] return selfmin_count中文语料一般取 5 到 10英文取 5 也行。低于这个频次的词训出来的向量基本是噪声——一个词只出现两次它的上下文信息不足以定位它的语义强行给它分配 100 维空间只会让邻近区域变乱。max_size是给超大语料兜底的词表十万以内其实可以不设。counts必须按同样的顺序保存下来负采样表要拿它算分布顺序错一位采样出来的全是错词损失照样在降向量质量却是坏的。3.2 词表与 ID 映射顺序就是 Embedding 的行号词表在内存里就是一个 list索引即 ID。这件事听着简单但有两个隐形约定必须写死第一id2token的顺序一旦生成就不再变动。后续增量训练、增量加词只能往尾部 append中间不能插。插一个词后面所有词的 ID 全体后移一位已经训好的 Embedding 矩阵全部错位效果比从头训还差。第二token2id和id2token必须同时序列化到vocab.json别只存前者然后靠遍历重建后者——dict 的遍历顺序在 Python 3.7 虽然稳定但依赖这个行为的人迟早会在某个版本切换上翻车。序列化的写法import json def save_vocab(vocab, path): with open(path, w, encodingutf-8) as f: json.dump({ id2token: vocab.id2token, # 顺序 Embedding 行号 counts: vocab.counts, min_count: vocab.min_count, }, f, ensure_asciiFalse) def load_vocab(path): with open(path, encodingutf-8) as f: d json.load(f) v Vocab(min_countd[min_count]) v.id2token d[id2token] v.token2id {w: i for i, w in enumerate(v.id2token)} v.counts d[counts] return v加载时token2id是从id2token重建的这样即使 JSON 文件被人手工改过、两者不一致重建后也一定自洽。counts也要一起走因为它和词表严格同序是负采样表的唯一依据。把这几个字段放在同一个文件里就是为了杜绝「词表换了、采样表没换」这种不同步问题。推理阶段的未登录词OOV处理要提前定好策略。简单做法是查不到就返回None让调用方跳过讲究一点的做法是在词表里预留unk但这个 token 训出来的向量意义有限用它替代任意未见词只会引入噪声。我的习惯是 OOV 直接跳过因为词向量本质是查表查不到就说明这个词在这个语料的统计里不存在硬给它一个向量是自欺欺人。3.3 滑动窗口采样与批次组装采样这一步是 word2vec 的性能关键。最朴素的写法是双重循环Python 层面几十万次迭代训练时间全花在解释器上。正确做法是把(center, context)对攒起来按批吐import random def make_batches(token_ids, window, batch_size, seed42): rng random.Random(seed) centers, contexts [], [] for i, center in enumerate(token_ids): # 动态窗口每次只在 [1, window] 里抽一个实际窗口 # 让近距离上下文被采到的概率更高这是原论文的经验做法 w rng.randint(1, window) lo, hi max(0, i - w), min(len(token_ids), i w 1) for j in range(lo, hi): if j i: continue centers.append(center) contexts.append(token_ids[j]) if len(centers) batch_size: yield centers, contexts centers, contexts [], [] if centers: # 尾批别丢 yield centers, contextswindow是最大窗口半径取 5 意味着看中心词左右各 5 个词以内。中文语料常用 5英文常用 5 到 10。w rng.randint(1, window)这一行是动态窗口原文里叫 dynamic window距离近的词对语义贡献大让它有更大概率被采样到距离远的词也偶尔进样本帮模型学到更宽泛的话题相关性。如果写成固定窗口等于每对词的权重一样训出来的向量在近义任务上会钝一些。batch_size是攒多少对再吐一次。词向量训练不像分类任务那么吃 batch1024 到 8192 都有人用越大梯度越稳但显存占用越高。我一般从 4096 起调CPU 上直接 1024。注意token_ids是整篇文档拼起来的 ID 序列不要按行切分后各自滑窗——句间边界处会漏掉跨句的上下文而且短句多的时候尾批处理很麻烦。整篇拼起来、只在文档之间不跨窗是更省事的做法。真要在文档边界切断可以在文档之间插一个分隔 token让窗口自然收敛。这个生成器一次只在内存里存一个 batch 的对语料再大也不会把内存撑爆。缺点是每轮 epoch 都要重扫一遍整个序列想加速可以把token_ids预先转成 numpy 数组取窗口时用切片而不是 Python 循环能快 3 到 5 倍。4. 模型核心Skip-gram、负采样与手写训练循环拆解4.1 Skip-gram 和 CBOW 怎么选两套结构方向相反CBOW 用上下文预测中心词Skip-gram 用中心词预测上下文。实现的差别只是输入输出的角色对调但训出来的向量性质不一样。CBOW 训练速度快因为每个窗口只产生一个样本多个上下文词取平均后预测一个中心词一轮 epoch 的样本数约等于词数。代价是它对高频词的收敛更快对低频词则很不友好——一个生僻词作为上下文词出现时它的向量被平均进一堆词里梯度信号被稀释。Skip-gram 反过来每个窗口产生2 * window个样本训练慢几倍但低频词因为每次都当中心词被单独预测一次向量质量明显更好。语料量不大、或者领域术语多医疗、法律这类的场景我基本都是 Skip-gram 起步。class SkipGram(nn.Module): def __init__(self, vocab_size, emb_dim, sparseTrue): super().__init__() self.in_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) self.out_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) bound 0.5 / emb_dim self.in_emb.weight.data.uniform_(-bound, bound) self.out_emb.weight.data.uniform_(-bound, bound)emb_dim常用 100 到 300。100 维在中小语料上够用300 维需要至少几千万词的语料才喂得饱语料不够硬上大维度向量会稀疏得像随机噪声。sparseTrue这个参数很关键它让 Embedding 的梯度以稀疏张量形式回传后面配SparseAdam才能省显存。0.5 / emb_dim这个初始化范围来自原论文维度越大初始化越窄目的是防止初始状态下内积过大导致 sigmoid 饱和。自己拍脑袋写normal(0, 1)初始化训练前期损失会卡在 0.69 附近不动看起来像没学实际是梯度全被压平了。4.2 两张 Embedding 表的分工这是最容易讲错的一块。Skip-gram 里有两张表in_emb是中心词向量out_emb是上下文词向量两者形状相同但作用不同。in_emb是最终要导出的那份查相似词、做类比、喂下游都只用它。out_emb只在训练时充当「分类器权重」把中心词向量和上下文词向量的内积映射成一个打分判断这对词是否真的在语料里共现过。训练结束后out_emb通常直接丢掉。有人试过把两张表相加或拼接在部分任务上确实有微弱提升但收益远不如多训几轮稳定而且会让下游的向量维度翻倍不划算。需要留意的是两张表不能混用。用in_emb训练、用out_emb查相似词结果会很怪——相似词列表里高频虚词扎堆语义近邻反而不见了。原因是out_emb学的是「作为上下文出现时的判别方向」它衡量的不是词义相似度。4.3 负采样损失把 V 分类问题降成 K1 分类完整 softmax 要在每个样本上对全词表算一遍归一化词表十万就是十万次内积训练一轮的成本不可接受。负采样的思路是不要求模型给出全词表的概率分布只要求它区分「真实共现的词对」和「随机凑出来的词对」。import torch import torch.nn as nn import torch.nn.functional as F class SkipGramNS(nn.Module): def __init__(self, vocab_size, emb_dim, neg_k5, sparseTrue): super().__init__() self.neg_k neg_k self.in_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) self.out_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) bound 0.5 / emb_dim self.in_emb.weight.data.uniform_(-bound, bound) self.out_emb.weight.data.uniform_(-bound, bound) def forward(self, center, context, neg): v self.in_emb(center) # [B, D] u_pos self.out_emb(context) # [B, D] u_neg self.out_emb(neg) # [B, K, D] pos_score torch.sum(v * u_pos, dim1) # [B] neg_score torch.bmm(u_neg, v.unsqueeze(2)).squeeze(2) # [B, K] # 正样本要 sigmoid 越大越好负样本要越小越好 loss -(F.logsigmoid(pos_score).mean() F.logsigmoid(-neg_score).sum(dim1).mean()) return lossneg_k是每个正样本配几个负样本小语料取 5 到 10大语料取 2 到 5。取大了每步计算量线性上涨收益递减。F.logsigmoid比先sigmoid再log数值稳输入在 -30 这种极端值时前者不会返回-inf后者会然后整个 loss 变nan。torch.bmm(u_neg, v.unsqueeze(2))这行是批量矩阵乘把[B, K, D]和[B, D, 1]乘成[B, K, 1]一次算完一个 batch 所有负样本的分数。写成双层循环也能跑但速度差一个数量级。负样本不是均匀抽的按词频的 0.75 次方抽import numpy as np def build_neg_table(counts, power0.75, table_size10_000_000): freq np.asarray(counts, dtypenp.float64) ** power prob freq / freq.sum() # 预先抽好一张大表训练时按下标随机取比每步调 np.random.choice 快得多 return np.random.choice(len(counts), sizetable_size, pprob).astype(np.int64) def sample_negatives(neg_table, batch_size, k, device): idx np.random.randint(0, len(neg_table), size(batch_size, k)) return torch.as_tensor(neg_table[idx], dtypetorch.long, devicedevice)0.75 次方这个指数的效果是压高频、抬低频一个出现一百万次的词权重被压到原来的万分之一量级出现十次的词权重相对被抬高。如果按原始频率抽负样本那负样本里全是「的」「是」模型学会的只是「区分实词和虚词」学不到实词之间的语义差异。table_size取一千万内存占用 80MB 左右是采样速度和精度的平衡点语料只有几十万词时降到一百万也够。4.4 手写训练循环与优化器选择optimizer 这块有明确的坑。词表十万、维度 100两张 Embedding 表就是两千万参数用普通Adam会给每个参数存两份动量显存多出 160MB 起步词表百万级时这个数字变成几个 G直接 OOM。def train(model, batches, neg_table, device, lr2e-3, epochs5, log_every200): model.to(device) # sparseTrue 的 Embedding 必须配 SparseAdam普通 Adam 只更新当步碰到的行 # 没被碰到的行动量不衰减训练久了会漂 optimizer torch.optim.SparseAdam(model.parameters(), lrlr) step 0 for epoch in range(epochs): model.train() running 0.0 for center, context in batches: center torch.as_tensor(center, dtypetorch.long, devicedevice) context torch.as_tensor(context, dtypetorch.long, devicedevice) neg sample_negatives(neg_table, len(center), model.neg_k, device) loss model(center, context, neg) optimizer.zero_grad(set_to_noneTrue) loss.backward() optimizer.step() running loss.item() step 1 if step % log_every 0: print(fepoch {epoch} step {step} loss {running / log_every:.4f}) running 0.0 yield epochSparseAdam只对稀疏梯度里出现过的行做更新没出现的行原样不动显存只和当步的稀疏度有关。代价是它不支持权重衰减、不支持梯度裁剪想加这些得自己手动做。zero_grad(set_to_noneTrue)比默认的置零快一点因为它把梯度设成None而不是填零张量省一次写内存。lr2e-3是SparseAdam在词向量任务上的常用起点。注意这个值和 batch_size 相关batch 从 1024 涨到 8192学习率可以适当放大到 5e-3但别线性放大词向量训练对学习率比分类任务敏感。损失从初始的 4 左右log(1) 5 * log(1)量级降到 2 到 2.5 之间属于正常降到 1 以下通常意味着过拟合或者负采样出了问题。yield epoch是为了让外层能按 epoch 存 checkpoint、跑验证。写成生成器而不是回调是因为训练中断后续跑、以及在 notebook 里逐 epoch 观察曲线这两种场景生成器的控制力更强。如果项目里带层次 Softmax 的可选实现思路是用词频建一棵 Huffman 树每个词对应一条从根到叶的路径损失变成路径上每个节点做一次二分类的累加复杂度从O(V)降到O(log V)。它和负采样二选一即可两个都开没有额外收益还会让实现复杂到不好调。语料里长尾词特别多的时候层次 Softmax 表现更稳词表本身不大时负采样实现简单、调参直观我更倾向后者。5. 排查手册训练不收敛、显存爆掉、相似词离谱的五个高频现场5.1 开跑前先过这三关第一关数据。wc -l data/processed/tokens.txt看一眼行数再head -3看看分词结果。见过太多次分词器把整段中文切成一个 token词表里就几百个词训练照样跑完、loss 照样降向量全是废的。分词后每行的平均 token 数应该在 10 以上低于这个数基本可以判定分词配置错了。第二关词表。打开vocab.json看看id2token前 20 个词。如果全是标点、数字占位符、单字说明min_count设低了或者清洗没做干净。正常语料的高频词应该是有实际语义的虚词和常用词。第三关显存/内存。词表大小乘以维度乘以 4float32 字节数再乘以 2两张表加上优化器状态算出静态占用。十万词表、100 维两张表约 80MB加 SparseAdam 的状态大概翻倍可控。百万词表、300 维就要 2.4GB 起先确认机器扛得住再开跑。注意先用--max_steps 500这类参数限制步数跑一遍确认日志、checkpoint、导出全链路通再放开跑完整训练。一个跑了两小时才发现导出路径写错的任务不值得。5.2 五个真实翻车现场现象一训练能跑完但查相似词返回的全是「的」「是」「了」。原因是高频词在损失里占的权重过大模型把大部分容量花在拟合这些词的分布上。检查两处负采样表是不是没做 0.75 次方降权以及有没有做高频词下采样。word2vec 原论文里有个 subsample 步骤按P(keep) (sqrt(f/t) 1) * t/f丢弃高频词t取 1e-5这一步能把「的」这类词的样本量砍掉九成以上。解决方式是在make_batches前对token_ids做一次按概率丢弃代码不到十行收益非常明显。现象二相似词列表里第一个永远是查询词自己或者近邻全是不相干的词。前者的原因是查询向量和矩阵里包含了自身算完没排除。后者多半是两张表用混了用out_emb做了查询。确认方式很简单打印model.in_emb.weight和model.out_emb.weight同一行看数值是否明显不同。修法是查询时统一从in_emb取并在排序结果里把查询词的 ID 剔除。def most_similar(emb_matrix, vocab, word, topk10): # 先 L2 归一化之后内积等价于余弦相似度 norm emb_matrix / (np.linalg.norm(emb_matrix, axis1, keepdimsTrue) 1e-9) idx vocab.token2id[word] scores norm norm[idx] order np.argsort(-scores) out [] for i in order: if i idx: # 必须跳过自身 continue out.append((vocab.id2token[i], float(scores[i]))) if len(out) topk: break return out现象三跑到几百步突然 OOM 或者 loss 变 nan。OOM 八成是优化器吃掉的。先用torch.optim.SparseAdam换掉Adam确认 Embedding 的sparseTrue没写漏。loss 变 nan 则是学习率太大导致logsigmoid在极端值上溢出把lr从 2e-3 降到 5e-4 先确认能稳定跑再逐步往上试。另外检查负采样表里有没有越界 ID——np.random.choice用的是len(counts)如果counts和词表长度不一致取出来的 ID 会超出 Embedding 行数PyTorch 报的是 index out of range但有些写法下会静默返回零向量损失看着正常向量是坏的。现象四中文语料训完查「手机」返回一堆「手」「机」「手机上」。分词粒度不一致。预处理用 jieba 切推理时用户输入的词是按别的方式切的或者语料里本身就混了不同粒度的结果。修法是把分词器和词表绑定保存vocab.json里存一份分词配置词典版本、是否用 HMM推理时必须用同一份配置。更彻底的做法是训完后做一次粒度合并把「手」这种明显是切分碎片的词从词表里踢掉代价是它们原本承担的上下文信息也一起丢了所以更好的是在预处理阶段就统一。现象五类比推理「A 对 B 相当于 C 对 ?」的结果完全随机。三个原因按可能性排序语料太小、向量没归一化、参与计算的词频太低。类比任务需要语料里存在足够多的平行结构比如「国王-男人-女人-女王」这类关系要在不同词对上反复出现几百万词的语料基本做不出稳定结果这不是代码问题。验证方式是换成相似词任务如果相似词正常而类比不行就是语料量不够别在代码里找原因。另外向量要先归一化再算a - b c不归一化的话模长大的词会主导结果返回的全是高频词。5.3 训练日志盯这三个数日志不用记太多三个数足够判断训练是不是健康。一是 loss 的滑动平均。正常形态是前 200 步快速下降之后进入缓慢下降的平台期。如果从头到尾平得像一条直线检查梯度是不是没回传sparseTrue配了非 SparseAdam 就会这样梯度不更新但不报错。如果抖动幅度超过均值的 30%说明 batch 太小或者学习率太大。二是每秒处理的样本对数。这个数突然掉一半通常是内存 swap 或者负采样表太大挤占了缓存。CPU 训练时尤其明显table_size从一千万降到一百万吞吐能回升不少。三是每隔若干步跑一次相似词抽查。选五个你确定语义关系的词看 top5 近邻。这一项比 loss 更能反映真实质量因为 loss 只说明模型拟合了共现统计不代表向量空间有好的几何结构。训练早期这些结果会很离谱中期开始收敛如果到后期还是乱的回头看数据。把这三次检查写进trainer.py的日志回调里按log_every触发别等训完再看。6. 二次开发从「能跑通」改成「能进生产」的三步接手这类项目多数人的需求不是重写是把它接到自己的流程里。这里给一条我实际走过的改造路线。第一步增量训练。新语料进来时不要从头训加载旧 checkpoint冻结词表、只对新词做随机初始化追加到尾部然后用小学习率原值的 1/5跑 1 到 2 个 epoch。关键是旧词 ID 一位都不能动新词只能 appendckpt torch.load(ckpt/w2v_epoch4.pt, map_locationcpu) vocab load_vocab(data/processed/vocab.json) model SkipGramNS(len(vocab.id2token), ckpt[emb_dim]) model.load_state_dict(ckpt[model]) # 新词追加到尾部旧 ID 保持不变这是增量训练唯一安全的方向 for w, c in new_counter.items(): if w not in vocab.token2id and c min_count: vocab.token2id[w] len(vocab.id2token) vocab.id2token.append(w) vocab.counts.append(c) # Embedding 行数变了需要重建模型或原地扩容然后才能继续训第二步导出给下游。下游最常见的是把词向量当成nn.Embedding的初始化权重这里有个必须对齐的地方vocab.json的id2token顺序和导出矩阵的行号必须完全一致padding_idx要单独指定别指望0号位置恰好是 pad 词。matrix torch.load(export/vectors.pt) # [V, D]顺序同 id2token embedding nn.Embedding.from_pretrained(matrix, freezeFalse, padding_idxvocab.token2id.get(pad, 0))第三步验收。改完之后别只看 loss跑三个指标随机抽 50 个词的 top10 近邻人工看语义合理性在一份小的词对相似度数据集上算 Spearman 相关系数和前一次训练的结果对比掉超过 0.03 就要回查最后在下游任务上跑一次验证集确认端到端指标没退化。这三个指标一起看能挡住绝大多数「loss 更低了但向量更差了」的假进步。最后一件事是我的个人习惯每次改采样逻辑或词表构建我都会先把改动前后各训一个 5 万步的小模型把两份相似词结果 diff 出来看。词向量这玩意儿不像分类任务有明确的 accuracy 兜底改动的影响经常要几千步之后才显形只看 loss 曲线等于黑匣子。从那以后我每次动预处理都强制走一遍这个小规模对比省下的返工时间远超那点训练开销。希望帮到你。本文还有配套的精品资源点击获取 
长河君 · 2026/10/26 22:01:13 · 阅读 1.2k PyTorch 手写 word2vec 词向量模型：从预处理到训练与增量部署
简介一份基于PyTorch实现的词向量模型代码包面向自然语言处理初学者以及希望动手理解Word Embedding的开发者。内容围绕CBOW、Skip-gram等经典思路展开演示如何把词汇映射为连续向量、构建词表并搭建简单网络完成训练可作为文本分类、情感分析等下游任务的预处理练习起点。压缩包共6个文件以3个py源码为主另有2个pyc编译缓存与1个txt英文语料源码负责数据处理与模型定义txt语料则提供含1803个单词的小型训练集整体仅约10KB便于本地快速跑通。已有3111人学习下载。篇幅虽小却覆盖数据预处理、模型构建、交叉熵损失计算与权重保存的完整链路并保留训练后的应用示例适合边读代码边调向量维度、窗口大小与迭代次数直观观察词向量质量变化再迁移到其他自然语言处理任务中。1. 从一份纯文本到一个能算距离的坐标PyTorch 词向量模型交付什么拿 PyTorch 手写一份词向量模型难的从来不是公式是让它在自己的语料上跑出像样的相似词。有个做评论情感分类的开发者标注数据只有三千条分类器怎么调都在 78% 徘徊后来他用几十万条无标注评论先训了一份词向量把句子里的词换成 100 维稠密向量再喂给分类器同一份标注数据准确率往上走了六个点。词向量模型干的不是业务是把「词」这种离散符号压成坐标让机器学习模型能算距离、做加减法。这份基于 PyTorch 实现的词向量模型走的是自然语言处理里最经典那条路纯文本输入Skip-gram 和 CBOW 两套结构带负采样与层次 Softmax训完能查相似词、做类比推理、导出可视化。想亲手把 word2vec 从公式写成代码的或者要拿自己领域语料重训一份向量的都合适词表从几万到几十万单卡甚至纯 CPU 都能跑。2. 项目骨架与技术选型七个 Python 文件撑起一次完整训练2.1 目录结构与各模块职责划分拆这类训练项目我第一眼看的是目录有没有把「数据加工」和「模型计算」分开。分不开的项目改一次分词规则就得动训练脚本后面基本没法维护。这份实现的分层是干净的word2vec-pytorch/ ├── configs/ │ └── sg_zh.yaml # 一次实验一份配置超参不写死在代码里 ├── data/ │ ├── raw/corpus.txt # 原始语料一行一段不做任何预处理 │ └── processed/ │ ├── tokens.txt # 分词结果词之间空格分隔 │ └── vocab.json # 词表与词频训练和推理共用同一份 ├── w2v/ │ ├── vocab.py # 词表构建、ID 映射、低频词截断 │ ├── dataset.py # 滑窗采样产出 (center, context) 批次 │ ├── model.py # SkipGram / CBOW 两个 nn.Module │ ├── trainer.py # 训练循环、日志、checkpoint 落盘 │ ├── evaluate.py # 相似词、类比推理、t-SNE 可视化 │ └── utils.py # 负采样表、随机种子、向量归一化 ├── scripts/ │ ├── preprocess.py # raw/corpus.txt - tokens.txt vocab.json │ ├── train.py # 入口读 yaml 开跑 │ └── export.py # 导出 .vec / .npy 给下游任务 ├── requirements.txt └── README.mdvocab.py是唯一持有「词 → ID」映射的地方训练和推理都从它拿避免了两边词表对不上的血泪事故。dataset.py只负责产出整数 ID不碰字符串这样后面换分词器不用改采样逻辑。model.py里两个 Module 都很薄核心就两张 Embedding 表加一个前向损失把复杂度压在损失函数而不是网络结构上这是 word2vec 这类模型的特点——网络浅但采样和损失讲究。trainer.py单独抽出来是为了能在不碰模型定义的前提下替换优化器、加梯度累积、加日志钩子。evaluate.py属于「训完必须马上能验证」的那部分很多实现把它漏掉结果训练跑完几小时想查个相似词还得现写脚本。configs/*.yaml这层看着多余实际是省后悔药的地方。词向量训练对一个参数极其敏感——窗口大小5 和 10 训出来的向量在类比任务上能差一截。把window、neg_k、min_count、subsample全部外置到配置文件一组实验一个文件比在代码里注释来注释去靠谱得多。2.2 依赖环境与最小可运行配置依赖不多锁死版本能省掉一堆玄学问题torch1.13,2.4 numpy1.21 tqdm pyyaml jieba # 中文语料才需要纯英文语料可以删掉 scikit-learn # 只给 t-SNE 可视化用不装也能训练Python 用 3.8 到 3.11 都行3.12 上部分老版本 torch 轮子会缺装之前先确认。CPU 机器直接pip install torch --index-url https://download.pytorch.org/whl/cpu别装默认包默认包会拖一个几百兆的 CUDA 运行时下来纯浪费。有卡的按 CUDA 版本对号入座装错了报的不是「版本不匹配」而是CUDA error: no kernel image is available for execution on the device这行报错能让人查半天。scikit-learn只在evaluate.py --plot时才 import写成函数内延迟导入这样在只有 torch 的环境里训练和查相似词照样能跑。jieba同理英文语料不需要它缺了也别让整个包 import 失败。提示先把data/raw/corpus.txt换成一份几千行的小语料跑通全流程再换大语料。大语料上出的问题和小语料上出的问题排查成本差十倍。2.3 一次训练从读取文本到落盘向量的完整链路整条链路是五步每一步的产物都能单独检查这点很重要scripts/preprocess.py读data/raw/corpus.txt做归一化和分词写出tokens.txt和vocab.json。w2v/vocab.py加载vocab.json把tokens.txt每行映射成 ID 序列。w2v/dataset.py在 ID 序列上滑窗攒够batch_size个(center, context)对就吐一批。w2v/trainer.py每个 batch 额外采neg_k个负样本前向、反向、更新按 epoch 存 checkpoint。scripts/export.py只导出in_emb输入向量表成.npy或文本格式供下游nn.Embedding.from_pretrained用。第 2 步和第 5 步是最容易出岔子的地方第 2 步决定词表顺序第 5 步决定下游能不能对上号。所以vocab.json里除了词表本身还要存一份id2token的数组顺序就是 Embedding 的行号导出时原样带出去。下游加载向量矩阵时行号必须和这份数组一一对应对不上就是「查出来的相似词毫无关联」这类问题且极难定位。链路中间产物tokens.txt建议保留别训完就删。语料里混进乱码、分词把「神经网络」切成「神经」「网络」这些问题都是回看这个文件才发现的。3. 语料预处理与词表构建从原始文本到 (center, context) 样本对3.1 分词、清洗与低频词截断预处理这步决定上限。词向量对语料质量极其敏感喂进去一堆 URL、表情符号、重复刷屏文本训出来的向量必然脏。import re import jieba def normalize(line: str) - str: line line.strip().lower() # 数字统一成占位符避免 num 之外的千奇百怪形态污染词表 line re.sub(r\d, num , line) # 只保留中英文、数字占位符和基本词字符其余一律替换成空格 line re.sub(r[^\w\u4e00-\u9fff], , line) return re.sub(r\s, , line).strip() def tokenize(line: str): # 英文按空格切中文交给 jieba实际项目里英文还要过一遍停用词 tokens [] for seg in line.split( ): if not seg: continue if re.search(r[\u4e00-\u9fff], seg): tokens.extend(w for w in jieba.lcut(seg) if w.strip()) else: tokens.append(seg) return tokensnormalize里lower()对中文没影响对英文是必须的否则Model和model会变成两个词各训一份向量两边都学不充分。数字替换成num是个取舍不替换的话年份、价格、ID 会撑爆词表且每个出现次数都很少替换之后这些小样本被归并成一个高频 token反而让上下文窗口更干净。tokenize里判断一段是不是中文再决定用不用 jieba是为了处理中英混排。纯英文语料整段走空格切分就够了比先 jieba 再拼回去快得多。停用词这块我的习惯是不在预处理阶段删而是在词表构建时通过min_count自然过滤——「的」「是」这类词频高删不掉但它们在负采样里会被降权后面会讲比硬删更符合 word2vec 的设计。低频词截断在vocab.py里做from collections import Counter class Vocab: def __init__(self, min_count5, max_sizeNone): self.min_count min_count # 词频低于此值的词直接丢弃 self.max_size max_size # 词表上限超出按词频截断 self.token2id {} self.id2token [] self.counts [] def build(self, token_lines): counter Counter() for tokens in token_lines: counter.update(tokens) items [(w, c) for w, c in counter.items() if c self.min_count] items.sort(keylambda x: -x[1]) # 按词频降序保证可复现 if self.max_size: items items[:self.max_size] self.id2token [w for w, _ in items] self.token2id {w: i for i, w in enumerate(self.id2token)} self.counts [c for _, c in items] return selfmin_count中文语料一般取 5 到 10英文取 5 也行。低于这个频次的词训出来的向量基本是噪声——一个词只出现两次它的上下文信息不足以定位它的语义强行给它分配 100 维空间只会让邻近区域变乱。max_size是给超大语料兜底的词表十万以内其实可以不设。counts必须按同样的顺序保存下来负采样表要拿它算分布顺序错一位采样出来的全是错词损失照样在降向量质量却是坏的。3.2 词表与 ID 映射顺序就是 Embedding 的行号词表在内存里就是一个 list索引即 ID。这件事听着简单但有两个隐形约定必须写死第一id2token的顺序一旦生成就不再变动。后续增量训练、增量加词只能往尾部 append中间不能插。插一个词后面所有词的 ID 全体后移一位已经训好的 Embedding 矩阵全部错位效果比从头训还差。第二token2id和id2token必须同时序列化到vocab.json别只存前者然后靠遍历重建后者——dict 的遍历顺序在 Python 3.7 虽然稳定但依赖这个行为的人迟早会在某个版本切换上翻车。序列化的写法import json def save_vocab(vocab, path): with open(path, w, encodingutf-8) as f: json.dump({ id2token: vocab.id2token, # 顺序 Embedding 行号 counts: vocab.counts, min_count: vocab.min_count, }, f, ensure_asciiFalse) def load_vocab(path): with open(path, encodingutf-8) as f: d json.load(f) v Vocab(min_countd[min_count]) v.id2token d[id2token] v.token2id {w: i for i, w in enumerate(v.id2token)} v.counts d[counts] return v加载时token2id是从id2token重建的这样即使 JSON 文件被人手工改过、两者不一致重建后也一定自洽。counts也要一起走因为它和词表严格同序是负采样表的唯一依据。把这几个字段放在同一个文件里就是为了杜绝「词表换了、采样表没换」这种不同步问题。推理阶段的未登录词OOV处理要提前定好策略。简单做法是查不到就返回None让调用方跳过讲究一点的做法是在词表里预留unk但这个 token 训出来的向量意义有限用它替代任意未见词只会引入噪声。我的习惯是 OOV 直接跳过因为词向量本质是查表查不到就说明这个词在这个语料的统计里不存在硬给它一个向量是自欺欺人。3.3 滑动窗口采样与批次组装采样这一步是 word2vec 的性能关键。最朴素的写法是双重循环Python 层面几十万次迭代训练时间全花在解释器上。正确做法是把(center, context)对攒起来按批吐import random def make_batches(token_ids, window, batch_size, seed42): rng random.Random(seed) centers, contexts [], [] for i, center in enumerate(token_ids): # 动态窗口每次只在 [1, window] 里抽一个实际窗口 # 让近距离上下文被采到的概率更高这是原论文的经验做法 w rng.randint(1, window) lo, hi max(0, i - w), min(len(token_ids), i w 1) for j in range(lo, hi): if j i: continue centers.append(center) contexts.append(token_ids[j]) if len(centers) batch_size: yield centers, contexts centers, contexts [], [] if centers: # 尾批别丢 yield centers, contextswindow是最大窗口半径取 5 意味着看中心词左右各 5 个词以内。中文语料常用 5英文常用 5 到 10。w rng.randint(1, window)这一行是动态窗口原文里叫 dynamic window距离近的词对语义贡献大让它有更大概率被采样到距离远的词也偶尔进样本帮模型学到更宽泛的话题相关性。如果写成固定窗口等于每对词的权重一样训出来的向量在近义任务上会钝一些。batch_size是攒多少对再吐一次。词向量训练不像分类任务那么吃 batch1024 到 8192 都有人用越大梯度越稳但显存占用越高。我一般从 4096 起调CPU 上直接 1024。注意token_ids是整篇文档拼起来的 ID 序列不要按行切分后各自滑窗——句间边界处会漏掉跨句的上下文而且短句多的时候尾批处理很麻烦。整篇拼起来、只在文档之间不跨窗是更省事的做法。真要在文档边界切断可以在文档之间插一个分隔 token让窗口自然收敛。这个生成器一次只在内存里存一个 batch 的对语料再大也不会把内存撑爆。缺点是每轮 epoch 都要重扫一遍整个序列想加速可以把token_ids预先转成 numpy 数组取窗口时用切片而不是 Python 循环能快 3 到 5 倍。4. 模型核心Skip-gram、负采样与手写训练循环拆解4.1 Skip-gram 和 CBOW 怎么选两套结构方向相反CBOW 用上下文预测中心词Skip-gram 用中心词预测上下文。实现的差别只是输入输出的角色对调但训出来的向量性质不一样。CBOW 训练速度快因为每个窗口只产生一个样本多个上下文词取平均后预测一个中心词一轮 epoch 的样本数约等于词数。代价是它对高频词的收敛更快对低频词则很不友好——一个生僻词作为上下文词出现时它的向量被平均进一堆词里梯度信号被稀释。Skip-gram 反过来每个窗口产生2 * window个样本训练慢几倍但低频词因为每次都当中心词被单独预测一次向量质量明显更好。语料量不大、或者领域术语多医疗、法律这类的场景我基本都是 Skip-gram 起步。class SkipGram(nn.Module): def __init__(self, vocab_size, emb_dim, sparseTrue): super().__init__() self.in_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) self.out_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) bound 0.5 / emb_dim self.in_emb.weight.data.uniform_(-bound, bound) self.out_emb.weight.data.uniform_(-bound, bound)emb_dim常用 100 到 300。100 维在中小语料上够用300 维需要至少几千万词的语料才喂得饱语料不够硬上大维度向量会稀疏得像随机噪声。sparseTrue这个参数很关键它让 Embedding 的梯度以稀疏张量形式回传后面配SparseAdam才能省显存。0.5 / emb_dim这个初始化范围来自原论文维度越大初始化越窄目的是防止初始状态下内积过大导致 sigmoid 饱和。自己拍脑袋写normal(0, 1)初始化训练前期损失会卡在 0.69 附近不动看起来像没学实际是梯度全被压平了。4.2 两张 Embedding 表的分工这是最容易讲错的一块。Skip-gram 里有两张表in_emb是中心词向量out_emb是上下文词向量两者形状相同但作用不同。in_emb是最终要导出的那份查相似词、做类比、喂下游都只用它。out_emb只在训练时充当「分类器权重」把中心词向量和上下文词向量的内积映射成一个打分判断这对词是否真的在语料里共现过。训练结束后out_emb通常直接丢掉。有人试过把两张表相加或拼接在部分任务上确实有微弱提升但收益远不如多训几轮稳定而且会让下游的向量维度翻倍不划算。需要留意的是两张表不能混用。用in_emb训练、用out_emb查相似词结果会很怪——相似词列表里高频虚词扎堆语义近邻反而不见了。原因是out_emb学的是「作为上下文出现时的判别方向」它衡量的不是词义相似度。4.3 负采样损失把 V 分类问题降成 K1 分类完整 softmax 要在每个样本上对全词表算一遍归一化词表十万就是十万次内积训练一轮的成本不可接受。负采样的思路是不要求模型给出全词表的概率分布只要求它区分「真实共现的词对」和「随机凑出来的词对」。import torch import torch.nn as nn import torch.nn.functional as F class SkipGramNS(nn.Module): def __init__(self, vocab_size, emb_dim, neg_k5, sparseTrue): super().__init__() self.neg_k neg_k self.in_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) self.out_emb nn.Embedding(vocab_size, emb_dim, sparsesparse) bound 0.5 / emb_dim self.in_emb.weight.data.uniform_(-bound, bound) self.out_emb.weight.data.uniform_(-bound, bound) def forward(self, center, context, neg): v self.in_emb(center) # [B, D] u_pos self.out_emb(context) # [B, D] u_neg self.out_emb(neg) # [B, K, D] pos_score torch.sum(v * u_pos, dim1) # [B] neg_score torch.bmm(u_neg, v.unsqueeze(2)).squeeze(2) # [B, K] # 正样本要 sigmoid 越大越好负样本要越小越好 loss -(F.logsigmoid(pos_score).mean() F.logsigmoid(-neg_score).sum(dim1).mean()) return lossneg_k是每个正样本配几个负样本小语料取 5 到 10大语料取 2 到 5。取大了每步计算量线性上涨收益递减。F.logsigmoid比先sigmoid再log数值稳输入在 -30 这种极端值时前者不会返回-inf后者会然后整个 loss 变nan。torch.bmm(u_neg, v.unsqueeze(2))这行是批量矩阵乘把[B, K, D]和[B, D, 1]乘成[B, K, 1]一次算完一个 batch 所有负样本的分数。写成双层循环也能跑但速度差一个数量级。负样本不是均匀抽的按词频的 0.75 次方抽import numpy as np def build_neg_table(counts, power0.75, table_size10_000_000): freq np.asarray(counts, dtypenp.float64) ** power prob freq / freq.sum() # 预先抽好一张大表训练时按下标随机取比每步调 np.random.choice 快得多 return np.random.choice(len(counts), sizetable_size, pprob).astype(np.int64) def sample_negatives(neg_table, batch_size, k, device): idx np.random.randint(0, len(neg_table), size(batch_size, k)) return torch.as_tensor(neg_table[idx], dtypetorch.long, devicedevice)0.75 次方这个指数的效果是压高频、抬低频一个出现一百万次的词权重被压到原来的万分之一量级出现十次的词权重相对被抬高。如果按原始频率抽负样本那负样本里全是「的」「是」模型学会的只是「区分实词和虚词」学不到实词之间的语义差异。table_size取一千万内存占用 80MB 左右是采样速度和精度的平衡点语料只有几十万词时降到一百万也够。4.4 手写训练循环与优化器选择optimizer 这块有明确的坑。词表十万、维度 100两张 Embedding 表就是两千万参数用普通Adam会给每个参数存两份动量显存多出 160MB 起步词表百万级时这个数字变成几个 G直接 OOM。def train(model, batches, neg_table, device, lr2e-3, epochs5, log_every200): model.to(device) # sparseTrue 的 Embedding 必须配 SparseAdam普通 Adam 只更新当步碰到的行 # 没被碰到的行动量不衰减训练久了会漂 optimizer torch.optim.SparseAdam(model.parameters(), lrlr) step 0 for epoch in range(epochs): model.train() running 0.0 for center, context in batches: center torch.as_tensor(center, dtypetorch.long, devicedevice) context torch.as_tensor(context, dtypetorch.long, devicedevice) neg sample_negatives(neg_table, len(center), model.neg_k, device) loss model(center, context, neg) optimizer.zero_grad(set_to_noneTrue) loss.backward() optimizer.step() running loss.item() step 1 if step % log_every 0: print(fepoch {epoch} step {step} loss {running / log_every:.4f}) running 0.0 yield epochSparseAdam只对稀疏梯度里出现过的行做更新没出现的行原样不动显存只和当步的稀疏度有关。代价是它不支持权重衰减、不支持梯度裁剪想加这些得自己手动做。zero_grad(set_to_noneTrue)比默认的置零快一点因为它把梯度设成None而不是填零张量省一次写内存。lr2e-3是SparseAdam在词向量任务上的常用起点。注意这个值和 batch_size 相关batch 从 1024 涨到 8192学习率可以适当放大到 5e-3但别线性放大词向量训练对学习率比分类任务敏感。损失从初始的 4 左右log(1) 5 * log(1)量级降到 2 到 2.5 之间属于正常降到 1 以下通常意味着过拟合或者负采样出了问题。yield epoch是为了让外层能按 epoch 存 checkpoint、跑验证。写成生成器而不是回调是因为训练中断后续跑、以及在 notebook 里逐 epoch 观察曲线这两种场景生成器的控制力更强。如果项目里带层次 Softmax 的可选实现思路是用词频建一棵 Huffman 树每个词对应一条从根到叶的路径损失变成路径上每个节点做一次二分类的累加复杂度从O(V)降到O(log V)。它和负采样二选一即可两个都开没有额外收益还会让实现复杂到不好调。语料里长尾词特别多的时候层次 Softmax 表现更稳词表本身不大时负采样实现简单、调参直观我更倾向后者。5. 排查手册训练不收敛、显存爆掉、相似词离谱的五个高频现场5.1 开跑前先过这三关第一关数据。wc -l data/processed/tokens.txt看一眼行数再head -3看看分词结果。见过太多次分词器把整段中文切成一个 token词表里就几百个词训练照样跑完、loss 照样降向量全是废的。分词后每行的平均 token 数应该在 10 以上低于这个数基本可以判定分词配置错了。第二关词表。打开vocab.json看看id2token前 20 个词。如果全是标点、数字占位符、单字说明min_count设低了或者清洗没做干净。正常语料的高频词应该是有实际语义的虚词和常用词。第三关显存/内存。词表大小乘以维度乘以 4float32 字节数再乘以 2两张表加上优化器状态算出静态占用。十万词表、100 维两张表约 80MB加 SparseAdam 的状态大概翻倍可控。百万词表、300 维就要 2.4GB 起先确认机器扛得住再开跑。注意先用--max_steps 500这类参数限制步数跑一遍确认日志、checkpoint、导出全链路通再放开跑完整训练。一个跑了两小时才发现导出路径写错的任务不值得。5.2 五个真实翻车现场现象一训练能跑完但查相似词返回的全是「的」「是」「了」。原因是高频词在损失里占的权重过大模型把大部分容量花在拟合这些词的分布上。检查两处负采样表是不是没做 0.75 次方降权以及有没有做高频词下采样。word2vec 原论文里有个 subsample 步骤按P(keep) (sqrt(f/t) 1) * t/f丢弃高频词t取 1e-5这一步能把「的」这类词的样本量砍掉九成以上。解决方式是在make_batches前对token_ids做一次按概率丢弃代码不到十行收益非常明显。现象二相似词列表里第一个永远是查询词自己或者近邻全是不相干的词。前者的原因是查询向量和矩阵里包含了自身算完没排除。后者多半是两张表用混了用out_emb做了查询。确认方式很简单打印model.in_emb.weight和model.out_emb.weight同一行看数值是否明显不同。修法是查询时统一从in_emb取并在排序结果里把查询词的 ID 剔除。def most_similar(emb_matrix, vocab, word, topk10): # 先 L2 归一化之后内积等价于余弦相似度 norm emb_matrix / (np.linalg.norm(emb_matrix, axis1, keepdimsTrue) 1e-9) idx vocab.token2id[word] scores norm norm[idx] order np.argsort(-scores) out [] for i in order: if i idx: # 必须跳过自身 continue out.append((vocab.id2token[i], float(scores[i]))) if len(out) topk: break return out现象三跑到几百步突然 OOM 或者 loss 变 nan。OOM 八成是优化器吃掉的。先用torch.optim.SparseAdam换掉Adam确认 Embedding 的sparseTrue没写漏。loss 变 nan 则是学习率太大导致logsigmoid在极端值上溢出把lr从 2e-3 降到 5e-4 先确认能稳定跑再逐步往上试。另外检查负采样表里有没有越界 ID——np.random.choice用的是len(counts)如果counts和词表长度不一致取出来的 ID 会超出 Embedding 行数PyTorch 报的是 index out of range但有些写法下会静默返回零向量损失看着正常向量是坏的。现象四中文语料训完查「手机」返回一堆「手」「机」「手机上」。分词粒度不一致。预处理用 jieba 切推理时用户输入的词是按别的方式切的或者语料里本身就混了不同粒度的结果。修法是把分词器和词表绑定保存vocab.json里存一份分词配置词典版本、是否用 HMM推理时必须用同一份配置。更彻底的做法是训完后做一次粒度合并把「手」这种明显是切分碎片的词从词表里踢掉代价是它们原本承担的上下文信息也一起丢了所以更好的是在预处理阶段就统一。现象五类比推理「A 对 B 相当于 C 对 ?」的结果完全随机。三个原因按可能性排序语料太小、向量没归一化、参与计算的词频太低。类比任务需要语料里存在足够多的平行结构比如「国王-男人-女人-女王」这类关系要在不同词对上反复出现几百万词的语料基本做不出稳定结果这不是代码问题。验证方式是换成相似词任务如果相似词正常而类比不行就是语料量不够别在代码里找原因。另外向量要先归一化再算a - b c不归一化的话模长大的词会主导结果返回的全是高频词。5.3 训练日志盯这三个数日志不用记太多三个数足够判断训练是不是健康。一是 loss 的滑动平均。正常形态是前 200 步快速下降之后进入缓慢下降的平台期。如果从头到尾平得像一条直线检查梯度是不是没回传sparseTrue配了非 SparseAdam 就会这样梯度不更新但不报错。如果抖动幅度超过均值的 30%说明 batch 太小或者学习率太大。二是每秒处理的样本对数。这个数突然掉一半通常是内存 swap 或者负采样表太大挤占了缓存。CPU 训练时尤其明显table_size从一千万降到一百万吞吐能回升不少。三是每隔若干步跑一次相似词抽查。选五个你确定语义关系的词看 top5 近邻。这一项比 loss 更能反映真实质量因为 loss 只说明模型拟合了共现统计不代表向量空间有好的几何结构。训练早期这些结果会很离谱中期开始收敛如果到后期还是乱的回头看数据。把这三次检查写进trainer.py的日志回调里按log_every触发别等训完再看。6. 二次开发从「能跑通」改成「能进生产」的三步接手这类项目多数人的需求不是重写是把它接到自己的流程里。这里给一条我实际走过的改造路线。第一步增量训练。新语料进来时不要从头训加载旧 checkpoint冻结词表、只对新词做随机初始化追加到尾部然后用小学习率原值的 1/5跑 1 到 2 个 epoch。关键是旧词 ID 一位都不能动新词只能 appendckpt torch.load(ckpt/w2v_epoch4.pt, map_locationcpu) vocab load_vocab(data/processed/vocab.json) model SkipGramNS(len(vocab.id2token), ckpt[emb_dim]) model.load_state_dict(ckpt[model]) # 新词追加到尾部旧 ID 保持不变这是增量训练唯一安全的方向 for w, c in new_counter.items(): if w not in vocab.token2id and c min_count: vocab.token2id[w] len(vocab.id2token) vocab.id2token.append(w) vocab.counts.append(c) # Embedding 行数变了需要重建模型或原地扩容然后才能继续训第二步导出给下游。下游最常见的是把词向量当成nn.Embedding的初始化权重这里有个必须对齐的地方vocab.json的id2token顺序和导出矩阵的行号必须完全一致padding_idx要单独指定别指望0号位置恰好是 pad 词。matrix torch.load(export/vectors.pt) # [V, D]顺序同 id2token embedding nn.Embedding.from_pretrained(matrix, freezeFalse, padding_idxvocab.token2id.get(pad, 0))第三步验收。改完之后别只看 loss跑三个指标随机抽 50 个词的 top10 近邻人工看语义合理性在一份小的词对相似度数据集上算 Spearman 相关系数和前一次训练的结果对比掉超过 0.03 就要回查最后在下游任务上跑一次验证集确认端到端指标没退化。这三个指标一起看能挡住绝大多数「loss 更低了但向量更差了」的假进步。最后一件事是我的个人习惯每次改采样逻辑或词表构建我都会先把改动前后各训一个 5 万步的小模型把两份相似词结果 diff 出来看。词向量这玩意儿不像分类任务有明确的 accuracy 兜底改动的影响经常要几千步之后才显形只看 loss 曲线等于黑匣子。从那以后我每次动预处理都强制走一遍这个小规模对比省下的返工时间远超那点训练开销。希望帮到你。本文还有配套的精品资源点击获取