pytorch实现基于Word2Vec的词嵌入

PyTorch 实现 Word2Vec（Skip-gram 模型） 的完整代码，使用 中文语料 进行训练，包括数据预处理、模型定义、训练和测试。

1. 主要特点

支持中文数据，基于 jieba 进行分词
使用 Skip-gram 进行训练，适用于小数据集
支持负采样，提升训练效率
使用 cosine similarity 计算相似单词

完整代码：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
import jieba
from collections import Counter
from sklearn.metrics.pairwise import cosine_similarity# ========== 1. 数据预处理 ==========
corpus = ["我们 喜欢 深度 学习","自然 语言 处理 是 有趣 的","人工智能 改变 了 世界","深度 学习 是 人工智能 的 重要 组成部分"
]# 超参数
window_size = 2      # 窗口大小
embedding_dim = 10   # 词向量维度
num_epochs = 100     # 训练轮数
learning_rate = 0.01 # 学习率
batch_size = 4       # 批大小
neg_samples = 5      # 负采样个数# 分词 & 构建词汇表
tokenized_corpus = [list(jieba.cut(sentence)) for sentence in corpus]
vocab = set(word for sentence in tokenized_corpus for word in sentence)
word2idx = {word: idx for idx, word in enumerate(vocab)}
idx2word = {idx: word for word, idx in word2idx.items()}# 统计词频
word_counts = Counter([word for sentence in tokenized_corpus for word in sentence])
total_words = sum(word_counts.values())# 计算负采样概率
word_freqs = {word: count / total_words for word, count in word_counts.items()}
word_powers = {word: freq ** 0.75 for word, freq in word_freqs.items()}
Z = sum(word_powers.values())
word_distribution = {word: prob / Z for word, prob in word_powers.items()}# 负采样函数
def negative_sampling(positive_word, num_samples=5):words = list(word_distribution.keys())probabilities = list(word_distribution.values())negatives = []while len(negatives) < num_samples:neg = np.random.choice(words, p=probabilities)if neg != positive_word:negatives.append(neg)return negatives# 生成 Skip-gram 训练数据
data = []
for sentence in tokenized_corpus:indices = [word2idx[word] for word in sentence]for center_idx in range(len(indices)):center_word = indices[center_idx]for offset in range(-window_size, window_size + 1):context_idx = center_idx + offsetif 0 <= context_idx < len(indices) and context_idx != center_idx:context_word = indices[context_idx]data.append((center_word, context_word))# 转换为 PyTorch 张量
data = [(torch.tensor(center), torch.tensor(context)) for center, context in data]# ========== 2. 定义 Word2Vec (Skip-gram) 模型 ==========
class Word2Vec(nn.Module):def __init__(self, vocab_size, embedding_dim):super(Word2Vec, self).__init__()self.embedding = nn.Embedding(vocab_size, embedding_dim)self.output_layer = nn.Linear(embedding_dim, vocab_size)def forward(self, center_word):embed = self.embedding(center_word)  # 获取中心词向量out = self.output_layer(embed)       # 计算词分布return out# 初始化模型
model = Word2Vec(len(vocab), embedding_dim)# ========== 3. 训练 Word2Vec ==========
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)for epoch in range(num_epochs):total_loss = 0random.shuffle(data)  # 每轮打乱数据for center_word, context_word in data:optimizer.zero_grad()output = model(center_word.unsqueeze(0))  # 预测词分布loss = criterion(output, context_word.unsqueeze(0))  # 计算损失loss.backward()optimizer.step()total_loss += loss.item()if (epoch + 1) % 10 == 0:print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss:.4f}")# ========== 4. 测试词向量 ==========
word_vectors = model.embedding.weight.data.numpy()# 计算单词相似度
def most_similar(word, top_n=3):if word not in word2idx:return "单词不在词汇表中"word_vec = word_vectors[word2idx[word]].reshape(1, -1)similarities = cosine_similarity(word_vec, word_vectors)[0]# 获取相似度最高的 top_n 个单词（排除自身）similar_idx = similarities.argsort()[::-1][1:top_n+1]return [(idx2word[idx], similarities[idx]) for idx in similar_idx]# 测试相似词
test_words = ["深度", "学习", "人工智能"]
for word in test_words:print(f"【{word}】的相似单词:", most_similar(word))