跳到主要内容

PyTorch Word2Vec 实现

介绍​

Word2Vec 是一种广泛使用的词嵌入技术，它能够将单词映射到连续的向量空间中，从而捕捉单词之间的语义关系。通过 Word2Vec，我们可以将单词表示为低维向量，这些向量能够反映单词的上下文关系，例如“国王” - “男人” + “女人” ≈ “女王”。

在本教程中，我们将使用 PyTorch 实现一个简单的 Word2Vec 模型，并解释其背后的原理。我们将从数据预处理开始，逐步构建模型，并最终训练它来生成词向量。

Word2Vec 的基本概念​

Word2Vec 有两种主要的模型架构：

  1. CBOW (Continuous Bag of Words): 通过上下文单词预测目标单词。
  2. Skip-Gram: 通过目标单词预测上下文单词。

在本教程中，我们将实现 Skip-Gram 模型。

Skip-Gram 模型​

Skip-Gram 模型的目标是通过一个单词来预测其周围的上下文单词。例如，给定句子 "I love natural language processing"，如果我们选择 "love" 作为目标单词，那么模型的目标是预测其周围的单词 "I" 和 "natural"。

数据预处理​

在实现 Word2Vec 之前，我们需要对文本数据进行预处理。这包括分词、构建词汇表以及生成训练样本。

python
import torch
from torch.utils.data import Dataset, DataLoader
from collections import Counter
import numpy as np

# 示例文本
text = "I love natural language processing and I love machine learning"

# 分词
words = text.lower().split()

# 构建词汇表
vocab = Counter(words)
vocab = sorted(vocab, key=vocab.get, reverse=True)
vocab_size = len(vocab)
word_to_idx = {word: i for i, word in enumerate(vocab)}
idx_to_word = {i: word for i, word in enumerate(vocab)}

# 生成训练样本
window_size = 2
data = []
for i, word in enumerate(words):
    for j in range(max(0, i - window_size), min(len(words), i + window_size + 1)):
        if i != j:
            data.append((word_to_idx[word], word_to_idx[words[j]]))

# 转换为 PyTorch 张量
data = torch.tensor(data, dtype=torch.long)

构建 Word2Vec 模型​

接下来，我们将构建一个简单的 Skip-Gram 模型。模型的核心是一个嵌入层，它将单词索引映射到低维向量空间。

python
import torch.nn as nn
import torch.optim as optim

class SkipGram(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(SkipGram, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.linear = nn.Linear(embedding_dim, vocab_size)
    
    def forward(self, x):
        embeds = self.embeddings(x)
        out = self.linear(embeds)
        return out

# 超参数
embedding_dim = 10
learning_rate = 0.001
num_epochs = 100

# 初始化模型、损失函数和优化器
model = SkipGram(vocab_size, embedding_dim)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

训练模型​

现在我们可以开始训练模型了。我们将使用随机梯度下降（SGD）来优化模型参数。

python
for epoch in range(num_epochs):
    total_loss = 0
    for target, context in data:
        # 前向传播
        output = model(target)
        loss = criterion(output, context)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(data):.4f}')

实际应用场景​

Word2Vec 在许多自然语言处理任务中都有广泛的应用，例如：

  • 文本分类: 通过词向量表示文本，可以提高分类模型的性能。
  • 机器翻译: 词向量可以帮助模型理解不同语言之间的语义关系。
  • 推荐系统: 通过分析用户评论中的词向量，可以更好地理解用户的偏好。

总结​

在本教程中，我们使用 PyTorch 实现了一个简单的 Word2Vec 模型。我们从数据预处理开始，逐步构建了 Skip-Gram 模型，并训练了模型来生成词向量。通过这个过程，我们了解了 Word2Vec 的基本概念及其在自然语言处理中的应用。

附加资源与练习​

提示

如果你对 Word2Vec 的更多细节感兴趣，可以深入研究 CBOW 模型，并尝试实现它。