Pytorch实现word2vec(Skip-gram训练方式)

示例代码和语料库来自于博客：（宝藏博主，其它博客对学习NLP很有用）

https://wmathor.com/index.php/archives/1443/

https://wmathor.com/index.php/archives/1435/

语料库下载地址：https://pan.baidu.com/s/10Bd3JxCCFTjBPNt0YROvZA 提取码：81fo

 

简易版本的word2vec实现

skip-gram原理简述

skip-gram是word2vec的一种训练方法，是核心思想是用中心词预测周围词，相比起用周围词预测中心词的CBOW训练方法，skip-gram训练的“难度更大”，因此训练出来的词向量往往也要比CBOW的要好一些。

从理论到代码最重要的一步就是要认识到在用中心词预测周围词的时候，比如当我们指定窗口为2，那么左右的周围词共有四个，skip-gram的训练过程不是一次性用中心词预测四个词，而是中心词和一个周围词组成一个训练样本，有4个周围词的话就有4个样本，即[中心词，周围词1]、[中心词，周围词2]...

在这个简易的实现版本中，只有一个相关的词向量矩阵，没有中心词矩阵和周围词矩阵，也不会涉及层次softmax或负采样这些优化措施，只是展现了最核心的一些原理。

 

导包

import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import torch.utils.data as Data

dtype = torch.FloatTensor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

 

文本预处理

这里没有使用大规模的语料库，只是使用了几句话，虽然数据很简单，但是在后面还是准备了相关的Dataset和Dataloader，在大规模语料库中，代码还是可以复用的。

# 文本预处理
sentences = ["jack like dog", "jack like cat", "jack like animal",
             "dog cat animal", "banana apple cat dog like", "dog fish milk like",
             "dog cat animal like", "jack like apple", "apple like", "jack like banana",
             "apple banana jack movie book music like", "cat dog hate", "cat dog like"]

word_sequence = " ".join(sentences).split()  # ['jack', 'like', 'dog', 'jack', 'like', 'cat', 'animal',...]
vocab = list(set(word_sequence))  # build words vocabulary，去重
word2idx = {w: i for i, w in enumerate(vocab)}  # {'apple': 0, 'fish': 1,..., }，注意，不固定！！！不一定apple对应的就是0，在真实的源码中，是按照词频来排序、分配序号的。

 

模型的相关参数定义

# 模型的相关参数
batch_size = 8
embedding_size = 2  # 词向量的维度是2
C = 2  # window size，即左右各两个周围词
voc_size = len(vocab)  # 词典的大小

 

数据处理，构造数据集

注意：其实这里的中心词和周围词的构建不应该遍历下面代码中的word_sequence来实现，下面代码错误的把每句话首位相连，但是不同的话之前其实是没有上下文关系的，所以这里需要处理下！需要遍历每句话来处理，而不是遍历整个word_sequence，感兴趣的可以改写下代码。

# 数据预处理
skip_grams = []
print(word2idx)
for idx in range(C, len(word_sequence) - C):
    center = word2idx[word_sequence[idx]]  # 中心词

    context_idx = list(range(idx - C, idx)) + list(range(idx + 1, idx + C + 1))  # 中心词左边的2个词+中心词右边的两个词
    context = [word2idx[word_sequence[i]] for i in context_idx]
    for w in context:
        skip_grams.append([center, w])  # 中心词和每个周围词组成一个训练样本


def make_data(skip_grams):
    input_data = []
    output_data = []
    for i in range(len(skip_grams)):
        # input_data转换为one-hot形式，output_data合成一个list
        input_data.append(np.eye(voc_size)[skip_grams[i][0]])
        output_data.append(skip_grams[i][1])
    return input_data, output_data


print(skip_grams)
input_data, output_data = make_data(skip_grams)
print(input_data)
print(output_data)
input_data, output_data = torch.Tensor(input_data), torch.LongTensor(output_data)
dataset = Data.TensorDataset(input_data, output_data)
loader = Data.DataLoader(dataset, batch_size, True)
"""
skip_grams: [[10, 2],[9, 8], [11, 5], ..., [11, 7], [11, 10], [11, 0]]
input_data: [array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.]),...]
output_data: [2, 0, 2, 0, 0, 10, 0, 11, 10, 2, 11, 2, 2, 0, 2, 0, 0, 11, 0, 8, 11, 2, 8, 10, 2, 0, 10,...]
"""

 

构建/实例化模型

# 构建模型
class Word2Vec(nn.Module):
    def __init__(self):
        super(Word2Vec, self).__init__()
        # W：one-hot到词向量的hidden layer
        self.W = nn.Parameter(torch.randn(voc_size, embedding_size).type((dtype)))
        # V：输出层的参数
        self.V = nn.Parameter(torch.randn(embedding_size, voc_size).type((dtype)))

    def forward(self, X):
        # X : [batch_size, voc_size] one-hot
        # torch.mm only for 2 dim matrix, but torch.matmul can use to any dim
        hidden_layer = torch.matmul(X, self.W)  # hidden_layer : [batch_size, embedding_size]
        output_layer = torch.matmul(hidden_layer, self.V)  # output_layer : [batch_size, voc_size]
        return output_layer


model = Word2Vec().to(device)
criterion = nn.CrossEntropyLoss().to(device)  # 多分类，交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=1e-3)  # Adam优化算法

 

训练和测试（词向量可视化）

# 训练
for epoch in range(2000):
    for i, (batch_x, batch_y) in enumerate(loader):
        batch_x = batch_x.to(device)
        batch_y = batch_y.to(device)
        pred = model(batch_x)
        loss = criterion(pred, batch_y)
        if (epoch + 1) % 1000 == 0:
            print(epoch + 1, i, loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


# 将每个词在平面直角坐标系中标记出来，看看各个词