本文基于经典 Word2Vec 算法,完整拆解工业级 CBOW 词向量模型实现全流程:文本预处理、词汇表构建、模型定义、训练收敛、推理验证、词向量保存,严格遵循深度学习训练规范,全程规避新手高频踩坑点。所有代码均为可直接运行的完整源码,开箱复用。
一、背景与核心原理
1. 本文目标
在自然语言处理中,计算机无法直接识别文本单词,词向量(Word Embedding)是将离散文本转化为稠密数值向量的核心技术,是文本分类、情感分析、机器翻译、大模型语义理解等所有 NLP 任务的基石。
Word2Vec 作为最经典的词向量算法,包含两大核心结构,本文完整实现CBOW(Continuous Bag-of-Words) 模型:通过上下文单词预测中心词,完成词向量的无监督训练。
2. CBOW 核心逻辑与算法对比
CBOW 核心思想
以上下文窗口大小 = 2 为例:给定一句话 We use NLP in daily life,取中心词NLP,提取其前后各 2 个单词['We','use','in','daily']作为上下文,模型的训练目标就是:输入上下文,精准预测出中心词NLP。
训练完成后,模型的 Embedding 层权重,就是我们最终需要的词向量。
CBOW 与 Skip-gram 核心对比
| 算法 | 核心逻辑 | 适用场景 | 收敛速度 |
|---|---|---|---|
| CBOW | 上下文单词预测中心词 | 高频词、小数据集场景 | 更快 |
| Skip-gram | 中心词预测上下文单词 | 低频词、大数据集场景 | 稍慢 |
3. 模型实现的 "黄金法则"
在动手写代码之前,必须明确实现流程,规避新手高频踩坑点:文本预处理 → 词汇表与索引映射构建 → 训练样本(上下文 - 中心词)构建先完成模型定义与设备适配 → 再执行训练循环(前向传播→损失计算→反向传播→参数更新)推理阶段必须关闭梯度计算,严禁更新模型权重测试集 / 推理数据必须复用训练阶段的词汇表,严禁重新构建词向量提取必须基于训练完成的 Embedding 层权重,保证语义一致性本文所有代码严格遵循这一流程。
二、环境准备
1. 环境依赖
pip install torch numpy tqdm
2. 库导入与说明
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
import numpy as np
💡 细节解析
torch/torch.nn:PyTorch 核心框架,用于模型定义与训练tqdm:生成训练进度条,直观监控训练进度numpy:用于词向量的数值处理与本地保存- 全程兼容 CPU/GPU/MPS 设备,无需手动修改设备配置
三、文本预处理全流程
1. 超参数定义与语料加载
# 上下文窗口大小:中心词前后各取2个单词
CONTEXT_SIZE = 2
# 训练语料(可自行替换为任意英文文本,中文需先分词)
english_text = """Natural Language Processing (NLP) is a key branch of artificial intelligence.
It helps computers understand and process human language.
We use NLP in daily life, such as chat messages, news articles, and online reviews.
With basic NLP skills, we can analyze text, count words, and find hidden information easily.
It is easy to learn and very useful for beginners."""
# 文本分词(按空格切分)
english_text = english_text.split()
💡 细节解析
CONTEXT_SIZE = 2是 Word2Vec 的经典配置,可根据文本长度调整,窗口越大,词向量包含的语义信息越丰富- 中文文本需先通过 jieba 等分词工具完成分词,再执行后续流程,不可直接按空格切分
2. 词汇表与索引映射构建
神经网络只能处理数值数据,因此必须建立单词与数字索引的双向映射:
# 构建词汇表(去重)
vocab = set(english_text)
vocab_size = len(vocab)
# 单词→索引 映射(用于文本转数值)
word_to_index = {word:i for i,word in enumerate(vocab)}
# 索引→单词 映射(用于模型预测结果转回文本)
idx_to_word = {i:word for i,word in enumerate(vocab)}
💡 细节解析
set(english_text)对文本去重,得到唯一词汇表,词汇表大小决定了 Embedding 层的输入维度- 双向映射是 NLP 任务的基础操作,训练、推理、结果解析全流程必须复用同一套映射,严禁中途修改,否则会出现语义错位
3. 训练样本构建(CBOW 核心步骤)
遍历文本,为每个中心词匹配对应的上下文,生成(上下文, 中心词)的训练样本:
data = []
for i in range(CONTEXT_SIZE, len(english_text)-CONTEXT_SIZE):
# 左边2个词(i-2, i-1)
left_words = [english_text[i - 2], english_text[i - 1]]
# 右边2个词(i+1, i+2)
right_words = [english_text[i + 1], english_text[i + 2]]
# 上下文 = 左2 + 右2
context = left_words + right_words
# 中心词(模型要预测的目标)
target = english_text[i]
data.append((context, target))
💡 细节解析
- 循环范围
range(CONTEXT_SIZE, len(english_text)-CONTEXT_SIZE)保证了每个中心词的前后都有足够的上下文单词,不会出现索引越界 - 最终生成的每个样本,上下文固定为 4 个单词(窗口大小 2×2),目标为 1 个中心词,完全匹配 CBOW 的输入输出要求
4. 文本转张量工具函数
# 上下文单词转换为PyTorch张量索引
def make_context_vector(context, word_to_index):
idxs = [word_to_index[w] for w in context]
return torch.tensor(idxs, dtype=torch.long)
💡 细节解析
- 该函数将文本形式的上下文,通过
word_to_index映射转为数字索引,再转为 PyTorch 可处理的长整型张量 dtype=torch.long是nn.Embedding层的强制输入类型,类型不匹配会直接报错
四、CBOW 模型定义(PyTorch 实现)
1. 设备自动适配
# 自动选择设备:GPU优先,无GPU则用CPU/MPS(苹果芯片)
device ='cuda' if torch.cuda.is_available() else 'mps' \
if torch.backends.mps.is_available() else 'cpu'
💡 细节解析
- 无需手动修改代码,自动适配 Windows/Linux/macOS 全平台,GPU 训练可大幅提升训练速度
- 后续所有数据、模型都必须迁移到该设备上,否则会出现张量设备不匹配报错
2. CBOW 模型完整实现
class CBOW(nn.Module):
def __init__(self, vocab_size, embedding_dim):
super(CBOW, self).__init__()
# 词嵌入层:核心层,将单词索引转为词向量
self.embedding = nn.Embedding(vocab_size, embedding_dim)
# 全连接投影层:特征变换
self.proj = nn.Linear(embedding_dim, 128)
# 输出层:分类预测,输出维度等于词汇表大小
self.output = nn.Linear(128, vocab_size)
def forward(self, inputs):
# CBOW核心操作:对上下文词向量求和,压缩为一个特征向量
embeds = sum(self.embedding(inputs)).view(1, -1)
# 全连接层+ReLU激活函数
out = F.relu(self.proj(embeds))
# 输出层计算
out = self.output(out)
# 对数softmax,适配NLLLoss损失函数
nll_prob = F.log_softmax(out, dim=1)
return nll_prob
💡 细节解析
- Embedding 层:是整个模型的核心,训练完成后,该层的权重就是我们最终需要的词向量。输入维度为词汇表大小,输出维度为我们设定的词向量维度。
- 词向量求和 :是 CBOW 的标志性操作,将 4 个上下文单词的词向量相加,融合为一个整体特征向量,
view(1,-1)将张量调整为适配全连接层的二维形状。 - 激活函数与输出 :ReLU 激活函数引入非线性,提升模型拟合能力;
log_softmax将输出转为对数概率分布,适配后续的负对数似然损失函数。
3. 模型初始化
# 初始化模型:词汇表大小、词向量维度64(可自行调整为32/128/256)
model = CBOW(vocab_size, 64).to(device)
五、模型训练全流程
1. 优化器与损失函数定义
# Adam优化器,收敛速度快,适配小数据集
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 损失函数:负对数似然损失,适配多分类任务
loss_function = nn.NLLLoss()
# 存储每轮训练损失,用于后续收敛性分析
losses = []
💡 细节解析
- 学习率
lr=0.001是 Adam 优化器的经典配置,小数据集无需调整,大数据集可适当降低 - NLLLoss 与 log_softmax 是固定搭配,等价于交叉熵损失,是文本分类任务的标准损失函数
2. 训练循环
# 开启训练模式
model.train()
# 训练100轮,可根据数据集大小调整
for epoch in tqdm(range(100)):
total_loss = 0
# 遍历每个训练样本
for context, target in data:
# 上下文转张量,并迁移到指定设备
context_vector = make_context_vector(context, word_to_index).to(device)
# 目标中心词转张量,并迁移到指定设备
target = torch.tensor([word_to_index[target]]).to(device)
# 1. 前向传播:模型预测
train_predict = model(context_vector)
# 2. 损失计算
loss = loss_function(train_predict, target)
# 3. 梯度清零
optimizer.zero_grad()
# 4. 反向传播:计算梯度
loss.backward()
# 5. 参数更新
optimizer.step()
# 累计本轮总损失
total_loss += loss.item()
# 保存本轮损失
losses.append(total_loss)
# 打印本轮训练结果
print(f"Epoch [{epoch+1}/100], Loss: {total_loss:.4f}")
💡 细节解析
- 训练模式 :
model.train()开启模型训练模式,启用 Dropout、BatchNorm 等训练专属层,推理时必须切换为model.eval() - 梯度清零 :
optimizer.zero_grad()必须放在反向传播之前,否则梯度会累计,导致模型参数更新错误,是新手最高频踩坑点 - 设备迁移:所有输入数据、标签、模型必须在同一设备上,否则会直接报错
- 损失打印:训练过程中损失持续下降,说明模型正常收敛,若损失波动不下降,需检查学习率、数据构建是否正确
3. 训练结果展示
✅ 损失从 231 持续下降至 0.0087,模型完美收敛,训练有效。
六、模型推理与效果验证
1. 推理代码实现
# 测试上下文:对应原文的"We use NLP in daily"
context = ['We', 'use', 'in', 'daily' ]
# 上下文转张量并迁移设备
context_vector = make_context_vector(context, word_to_index).to(device)
# 开启评估模式
model.eval()
# 推理阶段关闭梯度计算,节省内存,避免误更新模型权重
with torch.no_grad():
predict = model(context_vector)
# 取概率最大的索引,即为预测的中心词索引
max_idx = predict.argmax(1).item()
# 打印预测结果
print(f"\n上下文: {context}")
print(f"模型预测的中心词: {idx_to_word[max_idx]}")
💡 细节解析
model.eval()关闭训练专属层,保证推理结果稳定with torch.no_grad()是推理阶段的强制规范,关闭梯度计算,既可以提升推理速度,也能完全避免模型权重被误更新argmax(1)取概率最高的类别索引,通过idx_to_word映射转回文本单词,完成最终预测
2. 推理结果展示
上下文: ['We', 'use', 'in', 'daily']
模型预测的中心词: NLP
✅ 模型完全根据上下文,精准预测出了原文中的中心词NLP,词向量训练效果达标,语义信息被完整学习。
七、词向量提取与本地保存
训练完成后,Embedding 层的权重就是我们最终需要的词向量,可提取保存用于后续 NLP 任务:
# 打印词向量权重矩阵
print('\nCBOW 词向量权重矩阵=', model.embedding.weight)
# 词向量权重转为numpy数组
W = model.embedding.weight.cpu().detach().numpy()
# 构建 单词-词向量 映射字典,可直接通过单词查询对应词向量
word_2_vec = {}
for word in word_to_index.keys():
word_2_vec[word] = W[word_to_index[word],:]
print('词向量构建完成!')
# 保存词向量到本地npz文件,后续可直接加载复用
np.savez('word2vec_cbow词向量.npz', embedding_matrix=W)
# 加载验证
data = np.load('word2vec_cbow词向量.npz')
print("保存的文件内容:", data.files)
💡 细节解析
.cpu().detach().numpy()是 PyTorch 张量转 numpy 数组的标准流程:先迁移到 CPU,再脱离计算图,最后转为 numpy 数组- 保存的词向量可直接用于文本分类、相似度计算、大模型 Embedding 层初始化等下游任务,无需重新训练
八、常见避坑指南
问题 1:NumPy 2.0 与 PyTorch 不兼容报错
RuntimeError: Numpy is not available
原因
PyTorch 暂未完全适配 NumPy 2.0 版本,仅稳定支持 NumPy 1.x 系列,导致张量转 numpy 数组、文件保存时报错。
解决方案(二选一)
-
降级 NumPy 到稳定兼容版本(推荐)
pip install numpy==1.26.4
-
若无需保存词向量,注释掉 numpy 相关代码,不影响模型训练与推理
问题 2:张量设备不匹配报错
原因
模型在 GPU 上,而输入数据在 CPU 上,未完成设备迁移。
解决方案
所有输入数据、标签都必须通过.to(device)迁移到与模型相同的设备上,本文代码已做完整适配。
问题 3:模型训练损失不下降
常见原因与解决方案
- 学习率设置不当:调大 / 调小学习率,推荐 0.001-0.01 区间
- 训练数据构建错误:检查上下文与中心词的对应关系,确保样本正确
- 梯度未清零:确保
optimizer.zero_grad()在反向传播之前执行 - 词向量维度设置过小:小数据集推荐 32-128 维度,大数据集可提升至 256-512 维度
问题 4:推理结果完全错误
原因
推理时使用了新的词汇表,或未复用训练时的word_to_index映射,导致单词索引错位。
解决方案
训练、推理全流程必须使用同一套词汇表与索引映射,严禁中途修改或重新构建。
九、总结
经过以上全流程处理,我们实现了:✅ 完整的 CBOW(Word2Vec)模型 PyTorch 实现,代码可直接复用✅ 文本预处理、词汇表构建、训练样本生成全流程规范操作✅ 模型正常收敛,精准完成上下文→中心词的预测任务✅ 训练完成的词向量可提取、保存,直接用于下游 NLP 任务✅ 全程规避新手高频踩坑点,严格遵循深度学习训练规范