Skip to content
W3S101

PyTorch - 词嵌入

PyTorch - 使用 nn.Embedding 和 Skip-Gram 实现现代词嵌入

Section titled “PyTorch - 使用 nn.Embedding 和 Skip-Gram 实现现代词嵌入”

在自然语言处理（NLP）中，词嵌入（word embeddings）是词语的密集向量表示，能够捕捉词语的语义信息。本章重点介绍如何使用现代 PyTorch 方法，理解并实现一个基础的词嵌入模型 Word2Vec，特别是其 Skip-Gram 变体。虽然 Word2Vec 本身最初是用 C 语言实现的，但 PyTorch 允许灵活高效地创建自定义模型。

Word2Vec 的 Skip-Gram 模型通过预测给定目标词（target word）的上下文词（surrounding words，即周围的词）来学习词嵌入。我们将使用 PyTorch 的 nn.Embedding 层来实现这一点，这是处理嵌入的标准方法。

Skip-Gram 的核心概念

Section titled “Skip-Gram 的核心概念”

  1. 目标词 (Target Word)： 我们希望学习其嵌入的输入词。

  2. 上下文词 (Context Words / 正样本 Positive Samples)： 在语料库中出现在目标词附近的词。

  3. 负样本 (Negative Samples)： 在目标词的上下文中未出现的词，从词汇表中随机采样。这对于高效训练至关重要。

  4. 两个嵌入层 (Two Embedding Layers)： 通常，Word2Vec 为每个词使用两组嵌入：一组作为目标词时使用 (u_embeddings)，另一组作为上下文词时使用 (v_embeddings)。通常，最终的词向量只使用 u_embeddings。

步骤 1：必要的导入

Section titled “步骤 1：必要的导入”

首先，我们导入必要的 PyTorch 模块。请注意，torch.autograd.Variable 不再需要，因为 tensors 可以直接跟踪梯度。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim # 用于优化器


# 示例：检查 CUDA 是否可用并设置设备
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

步骤 2：实现 Skip-Gram 模型

Section titled “步骤 2：实现 Skip-Gram 模型”

我们将定义一个继承自 nn.Module 的类 SkipGramModel。该模型将包含两个 nn.Embedding 层：一个用于目标词，一个用于上下文词。

class SkipGramModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(SkipGramModel, self).__init__()
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim


        # Embedding layer for target words (center words)
        self.u_embeddings = nn.Embedding(vocab_size, embedding_dim, sparse=True)
        # Embedding layer for context words
        self.v_embeddings = nn.Embedding(vocab_size, embedding_dim, sparse=True)


        self.init_embeddings()


    def init_embeddings(self):
        init_range = 0.5 / self.embedding_dim
        self.u_embeddings.weight.data.uniform_(-init_range, init_range)
        # Initialize v_embeddings similar to u_embeddings, not to zero
        self.v_embeddings.weight.data.uniform_(-init_range, init_range)


    def forward(self, target_word_indices, positive_context_indices, negative_context_indices):
        """
        target_word_indices: shape 为 (batch_size) 的 Tensor
        positive_context_indices: shape 为 (batch_size) 的 Tensor
        negative_context_indices: shape 为 (batch_size, num_negative_samples) 的 Tensor
        """
        # 获取目标词的嵌入
        # Shape: (batch_size, embedding_dim)
        emb_target = self.u_embeddings(target_word_indices)


        # 获取正上下文词的嵌入
        # Shape: (batch_size, embedding_dim)
        emb_positive_context = self.v_embeddings(positive_context_indices)


        # 计算正样本对的分数 (目标词 * 正上下文词)
        # 分数越高表示越相似
        # torch.mul 执行元素级乘法
        # torch.sum 沿 dim=1 对点积分量求和
        positive_score = torch.sum(torch.mul(emb_target, emb_positive_context), dim=1)
        positive_loss = F.logsigmoid(positive_score) # 最大化这个分数


        # 获取负上下文词的嵌入
        # Shape: (batch_size, num_negative_samples, embedding_dim)
        emb_negative_context = self.v_embeddings(negative_context_indices)


        # 计算负样本对的分数 (目标词 * 负上下文词)
        # torch.bmm (批量矩阵乘法) 在这里效率很高
        # 对于 bmm，emb_target 需要 unsqueeze 成 (batch_size, embedding_dim, 1)
        # 或者，更常见的是 (batch_size, 1, embedding_dim) 与 (batch_size, num_negative_samples, embedding_dim).transpose(1,2) 相乘
        # 更简单的方式：如果 num_negative_samples 较小，可以迭代或使用广播；或者小心地使用元素级乘法和求和。
        # 对于本例，在重塑目标嵌入后使用批量矩阵乘法。
        # emb_target_expanded shape: (batch_size, 1, embedding_dim)
        emb_target_expanded = emb_target.unsqueeze(1)


        # emb_negative_context shape: (batch_size, num_negative_samples, embedding_dim)
        # 我们想要计算每个负样本与目标嵌入的点积。
        # (batch_size, 1, embedding_dim) @ (batch_size, embedding_dim, num_negative_samples)
        # -> (batch_size, 1, num_negative_samples)
        negative_scores = torch.bmm(emb_target_expanded, emb_negative_context.transpose(1, 2)).squeeze(1)
        negative_loss = F.logsigmoid(-negative_scores) # 最小化这些分数（最大化 -scores）


        # 总损失是负对数似然之和
        # 我们想要最大化 positive_loss 和所有 negative_loss 元素的总和
        # 因此我们最小化 -(positive_loss 之和 + 所有 negative_loss 元素之和)
        total_loss = -(torch.sum(positive_loss) + torch.sum(negative_loss))
        return total_loss


    def get_word_embeddings(self):
        # 通常，u_embeddings 被用作最终的词向量
        return self.u_embeddings.weight.data.cpu().numpy()


    def save_embedding_vectors(self, id_to_word_map, output_file_path, use_cuda_device=False):
        if use_cuda_device:
            embeddings = self.u_embeddings.weight.cpu().data.numpy()
        else:
            embeddings = self.u_embeddings.weight.data.numpy()


        with open(output_file_path, 'w', encoding='utf-8') as fout:
            fout.write(f'{len(id_to_word_map)} {self.embedding_dim}\n')
            for word_id, word_string in id_to_word_map.items():
                embedding_vector = embeddings[word_id]
                embedding_str = ' '.join(map(str, embedding_vector))
                fout.write(f'{word_string} {embedding_str}\n')
        print(f"Embeddings saved to {output_file_path}")


# 示例用法（概念性的 - 未显示完整的训练循环）
def example_training_step(model, optimizer, target, pos_context, neg_context):
    model.train() # 设置模型为训练模式
    optimizer.zero_grad() # 清除之前的梯度
    loss = model(target, pos_context, neg_context)
    loss.backward() # 计算梯度
    optimizer.step() # 更新模型参数
    return loss.item()

步骤 3：使用模型（概念性）

Section titled “步骤 3：使用模型（概念性）”

要训练此模型，你需要一个大型文本语料库、一个词汇表（将词映射到整数 ID），以及一种生成 (target, positive_context, negative_context) 批量数据的方法。将使用 torch.optim.SGD 或 torch.optim.Adam 等优化器来更新模型的权重。

if __name__ == '__main__':
    # --- 超参数 (示例) ---
    VOCAB_SIZE = 10000  # 你的词汇表大小
    EMBEDDING_DIM = 100 # 期望的词向量维度
    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


    # --- 实例化模型 ---
    model = SkipGramModel(VOCAB_SIZE, EMBEDDING_DIM) # .to(device)
    # print(model)


    # --- 示例用虚拟数据（替换为实际数据加载）---
    # BATCH_SIZE = 64
    # NUM_NEGATIVE_SAMPLES = 5
    # dummy_target = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE,)) # .to(device)
    # dummy_pos_context = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE,)) # .to(device)
    # dummy_neg_context = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, NUM_NEGATIVE_SAMPLES)) # .to(device)


    # --- 优化器 ---
    # optimizer = optim.Adam(model.parameters(), lr=0.001)


    # --- 训练循环 (概念性) ---
    # for epoch in range(num_epochs):
    #     for target, pos_context, neg_context in data_loader:
    #         loss = example_training_step(model, optimizer, target, pos_context, neg_context)
    #         # print(f'Epoch {epoch}, Loss: {loss}')


    # --- 保存嵌入 ---
    # id_to_word = {i: str(i) for i in range(VOCAB_SIZE)} # 示例 id_to_word 映射
    # model.save_embedding_vectors(id_to_word, 'word_embeddings.txt')
    print("SkipGramModel 已定义。请填写数据加载和训练循环以使用。")

应用与进一步学习

Section titled “应用与进一步学习”

通过 Word2Vec 学习到的词嵌入是许多 NLP 任务的基础，包括：

  • 情感分析
  • 机器翻译
  • 文本分类
  • 问答系统

虽然从头开始实现 Word2Vec 是一个很好的学习练习，但在许多应用中，通常会使用预训练的词嵌入（如 GloVe、FastText，或来自 BERT 等 Transformer 模型）或作为更大、特定任务神经网络的一部分来微调 nn.Embedding 层。你可以在 PyTorch 官方文档中找到关于 nn.Embedding 和相关 PyTorch 功能的更多详细信息。