欢迎关注公众号:DeepL Newer

MENU

PyTorch实现Word2Vec

April 13, 2020 • Read: 6558 • Deep Learning阅读设置

本文主要是使用PyTorch复现word2vec论文

PyTorch中的nn.Embedding

实现关键是nn.Embedding()这个API,首先看一下它的参数说明

其中两个必选参数num_embeddings表示单词的总数目,embedding_dim表示每个单词需要用什么维度的向量表示。而nn.Embedding权重的维度也是(num_embeddings, embedding_dim),默认是随机初始化的

import torch
import torch.nn as nn

embeds = nn.Embedding(2, 5)
embeds.weight
# 输出:
Parameter containing:
tensor([[-1.1454,  0.3675, -0.3718,  0.3733,  0.5979],
        [-0.7952, -0.9794,  0.6292, -0.3633, -0.2037]], requires_grad=True)

如果使用预训练好的词向量,则采用

pretrained_weight = np.array(pretrained_weight)
embeds.weight.data.copy_(torch.from_numpy(pretrained_weight))

想要查看某个词的词向量,需要传入这个词在词典中的index,并且这个index得是LongTensor型的

embeds = nn.Embedding(100, 10)
embeds(torch.LongTensor([50]))
# 输出
tensor([[-1.9562e-03,  1.8971e+00,  7.0230e-01, -6.3762e-01, -1.9426e-01,
          3.4200e-01, -2.0908e+00, -3.0827e-01,  9.6250e-01, -7.2700e-01]],
       grad_fn=<EmbeddingBackward>)

过程详解

具体的word2vec理论可以在我的这篇博客看到,这里就不多赘述。下面说一下实现部分的细节

首先Embedding层输入的shape是[batchsize, seq_len],输出的shape是[batchsize, seq_len, embedding_dim]

上图的流程是把文章中的单词使用词向量来表示

  1. 提取文章所有的单词,把所有的单词按照频次降序排序(取前4999个,表示常出现的单词。其余所有单词均用'<UNK>'表示。所以一共有5000个单词)
  2. 5000个单词使用one-hot编码
  3. 通过训练会生成一个$5000\times 300$的矩阵,每一行向量表示一个词的词向量。这里的300是人为指定,想要每个词最终编码为词向量的维度,你也可以设置成别的

这个矩阵如何获得呢?在Skip-gram模型中,首先会随机初始化这个矩阵,然后通过一层神经网络来训练。最终这个一层神经网络的所有权重,就是要求的词向量的矩阵

从上面的图中看到,我们所学习的embedding层是一个训练任务的一小部分,根据任务目标反向传播,学习到embedding层里的权重weight。这个weight是类似一种字典的存在,他能根据你输入的one-hot向量查到相应的Embedding vector

Pytorch实现

导包

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as tud

from collections import Counter
import numpy as np
import random

import scipy
from sklearn.metrics.pairwise import cosine_similarity

random.seed(1)
np.random.seed(1)
torch.manual_seed(1)

C = 3 # context window
K = 15 # number of negative samples
epochs = 2
MAX_VOCAB_SIZE = 10000
EMBEDDING_SIZE = 100
batch_size = 32
lr = 0.2

上面的代码我想应该没有不明白的,C就是论文中选取左右多少个单词作为背景词。这里我使用的是负采样来近似训练,K=15表示随机选取15个噪声词。MAX_VOCAB_SIZE=10000表示这次实验训练10000个词的词向量,但实际上我只会选出语料库中出现次数最多的9999个词,还有一个词是<UNK>用来表示所有的其它词。每个词的词向量维度为EMBEDDING_SIZE

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

文件中的内容是英文文本,去除了标点符号,每个单词之间用空格隔开

读取文本数据并处理

with open('text8.train.txt') as f:
    text = f.read() # 得到文本内容

text = text.lower().split() # 分割成单词列表
vocab_dict = dict(Counter(text).most_common(MAX_VOCAB_SIZE - 1)) # 得到单词字典表,key是单词,value是次数
vocab_dict['<UNK>'] = len(text) - np.sum(list(vocab_dict.values())) # 把不常用的单词都编码为"<UNK>"
word2idx = {word:i for i, word in enumerate(vocab_dict.keys())}
idx2word = {i:word for i, word in enumerate(vocab_dict.keys())}
word_counts = np.array([count for count in vocab_dict.values()], dtype=np.float32)
word_freqs = word_counts / np.sum(word_counts)
word_freqs = word_freqs ** (3./4.)

最后一行代码,word_freqs存储了每个单词的频率,然后又将所有的频率变为原来的0.75次方,因为word2vec论文里面推荐这么做,当然你不改变这个值也行

实现DataLoader

接下来我们需要实现一个DataLoader,DataLoader可以帮助我们轻松打乱数据集,迭代的拿到一个mini-batch的数据等。一个DataLoader需要以下内容:

  1. 把所有word编码成数字
  2. 保存vocabulary,单词count、normalized word frequency
  3. 每个iteration sample一个中心词
  4. 根据当前的中心词返回context单词
  5. 根据中心词sample一些negative单词
  6. 返回sample出的所有数据

为了使用DataLoader,我们需要定义以下两个function

  • __len__():返回整个数据集有多少item
  • __getitem__(idx):根据给定的idx返回一个item

这里有一个好的tutorial介绍如何使用PyTorch DataLoader

class WordEmbeddingDataset(tud.Dataset):
    def __init__(self, text, word2idx, word_freqs):
        ''' text: a list of words, all text from the training dataset
            word2idx: the dictionary from word to index
            word_freqs: the frequency of each word
        '''
        super(WordEmbeddingDataset, self).__init__() # #通过父类初始化模型,然后重写两个方法
        self.text_encoded = [word2idx.get(word, word2idx['<UNK>']) for word in text] # 把单词数字化表示。如果不在词典中,也表示为unk
        self.text_encoded = torch.LongTensor(self.text_encoded) # nn.Embedding需要传入LongTensor类型
        self.word2idx = word2idx
        self.word_freqs = torch.Tensor(word_freqs)
        
        
    def __len__(self):
        return len(self.text_encoded) # 返回所有单词的总数,即item的总数
    
    def __getitem__(self, idx):
        ''' 这个function返回以下数据用于训练
            - 中心词
            - 这个单词附近的positive word
            - 随机采样的K个单词作为negative word
        '''
        center_words = self.text_encoded[idx] # 取得中心词
        pos_indices = list(range(idx - C, idx)) + list(range(idx + 1, idx + C + 1)) # 先取得中心左右各C个词的索引
        pos_indices = [i % len(self.text_encoded) for i in pos_indices] # 为了避免索引越界,所以进行取余处理
        pos_words = self.text_encoded[pos_indices] # tensor(list)
        
        neg_words = torch.multinomial(self.word_freqs, K * pos_words.shape[0], True)
        # torch.multinomial作用是对self.word_freqs做K * pos_words.shape[0]次取值,输出的是self.word_freqs对应的下标
        # 取样方式采用有放回的采样,并且self.word_freqs数值越大,取样概率越大
        # 每采样一个正确的单词(positive word),就采样K个错误的单词(negative word),pos_words.shape[0]是正确单词数量
        
        # while 循环是为了保证 neg_words中不能包含背景词
        while len(set(pos_indices.numpy().tolist()) & set(neg_words.numpy().tolist())) > 0:
            neg_words = torch.multinomial(self.word_freqs, K * pos_words.shape[0], True)

        return center_words, pos_words, neg_words

每一行代码详细的注释都写在上面了,其中有一行代码需要特别说明一下,就是注释了tensor(list)的那一行,因为text_encoded本身是个tensor,而传入的pos_indices是一个list。下面举个例子就很好理解这句代码的作用了

a = torch.tensor([2, 3, 3, 8, 4, 6, 7, 8, 1, 3, 5, 0], dtype=torch.long)
b = [2, 3, 5, 6]
print(a[b])
# tensor([3, 8, 6, 7])

通过下面两行代码即可得到DataLoader

dataset = WordEmbeddingDataset(text, word2idx, word_freqs)
dataloader = tud.DataLoader(dataset, batch_size, shuffle=True)

可以随便打印一下看看

next(iter(dataset))
'''
(tensor(4813),
 tensor([  50, 9999,  393, 3139,   11,    5]),
 tensor([  82,    0, 2835,   23,  328,   20, 2580, 6768,   34, 1493,   90,    5,
          110,  464, 5760, 5368, 3899, 5249,  776,  883, 8522, 4093,    1, 4159,
         5272, 2860, 9999,    6, 4880, 8803, 2778, 7997, 6381,  264, 2560,   32,
         7681, 6713,  818, 1219, 1750, 8437, 1611,   12,   42,   24,   22,  448,
         9999,   75, 2424, 9970, 1365, 5320,  878,   40, 2585,  790,   19, 2607,
            1,   18, 3847, 2135,  174, 3446,  191, 3648, 9717, 3346, 4974,   53,
          915,   80,   78, 6408, 4737, 4147, 1925, 4718,  737, 1628, 6160,  894,
         9373,   32,  572, 3064,    6,  943]))
'''

定义PyTorch模型

class EmbeddingModel(nn.Module):
    def __init__(self, vocab_size, embed_size):
        super(EmbeddingModel, self).__init__()
        
        self.vocab_size = vocab_size
        self.embed_size = embed_size
        
        self.in_embed = nn.Embedding(self.vocab_size, self.embed_size)
        self.out_embed = nn.Embedding(self.vocab_size, self.embed_size)
        
    def forward(self, input_labels, pos_labels, neg_labels):
        ''' input_labels: center words, [batch_size]
            pos_labels: positive words, [batch_size, (window_size * 2)]
            neg_labels:negative words, [batch_size, (window_size * 2 * K)]
            
            return: loss, [batch_size]
        '''
        input_embedding = self.in_embed(input_labels) # [batch_size, embed_size]
        pos_embedding = self.out_embed(pos_labels)# [batch_size, (window * 2), embed_size]
        neg_embedding = self.out_embed(neg_labels) # [batch_size, (window * 2 * K), embed_size]
        
        input_embedding = input_embedding.unsqueeze(2) # [batch_size, embed_size, 1]
        
        pos_dot = torch.bmm(pos_embedding, input_embedding) # [batch_size, (window * 2), 1]
        pos_dot = pos_dot.squeeze(2) # [batch_size, (window * 2)]
        
        neg_dot = torch.bmm(neg_embedding, -input_embedding) # [batch_size, (window * 2 * K), 1]
        neg_dot = neg_dot.squeeze(2) # batch_size, (window * 2 * K)]
        
        log_pos = F.logsigmoid(pos_dot).sum(1) # .sum()结果只为一个数,.sum(1)结果是一维的张量
        log_neg = F.logsigmoid(neg_dot).sum(1)
        
        loss = log_pos + log_neg
        
        return -loss
    
    def input_embedding(self):
        return self.in_embed.weight.detach().numpy()

model = EmbeddingModel(MAX_VOCAB_SIZE, EMBEDDING_SIZE)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

这里为什么要分两个embedding层来训练?很明显,对于任一一个词,它既有可能作为中心词出现,也有可能作为背景词出现,所以每个词需要用两个向量去表示。in_embed训练出来的权重就是每个词作为中心词的权重。out_embed训练出来的权重就是每个词作为背景词的权重。那么最后到底用什么向量来表示一个词呢?是中心词向量?还是背景词向量?按照Word2Vec论文所写,推荐使用中心词向量,所以这里我最后返回的是in_embed.weight。如果上面我说的你不太明白,可以看我之前的Word2Vec详解

bmm(a, b),batch matrix multiply。函数中的两个参数a,b都是维度为3的tensor,并且这两个tensor的第一个维度必须相同,后面两个维度必须满足矩阵乘法的要求

batch1 = torch.randn(10, 3, 4)
batch2 = torch.randn(10, 4, 5)
res = torch.bmm(batch1, batch2)
print(res.size())
# torch.Size([10, 3, 5])

训练模型

for e in range(1):
    for i, (input_labels, pos_labels, neg_labels) in enumerate(dataloader):
        input_labels = input_labels.long()
        pos_labels = pos_labels.long()
        neg_labels = neg_labels.long()

        optimizer.zero_grad()
        loss = model(input_labels, pos_labels, neg_labels).mean()
        loss.backward()

        optimizer.step()

        if i % 100 == 0:
            print('epoch', e, 'iteration', i, loss.item())

embedding_weights = model.input_embedding()
torch.save(model.state_dict(), "embedding-{}.th".format(EMBEDDING_SIZE))

如果没有GPU,训练时间可能比较长

词向量应用

我们可以写个函数,找出与某个词相近的一些词,比方说输入good,他能帮我找出nice,better,best之类的

def find_nearest(word):
    index = word2idx[word]
    embedding = embedding_weights[index]
    cos_dis = np.array([scipy.spatial.distance.cosine(e, embedding) for e in embedding_weights])
    return [idx2word[i] for i in cos_dis.argsort()[:10]]
for word in ["two", "america", "computer"]:
    print(word, find_nearest(word))
# 输出
two ['two', 'zero', 'four', 'one', 'six', 'five', 'three', 'nine', 'eight', 'seven']
america ['america', 'states', 'japan', 'china', 'usa', 'west', 'africa', 'italy', 'united', 'kingdom']
computer ['computer', 'machine', 'earth', 'pc', 'game', 'writing', 'board', 'result', 'code', 'website']

完整代码链接

nn.Linear VS. nn.Embedding

Word2Vec论文中给出的架构其实就一个单层神经网络,那么为什么不直接用nn.Linear()来训练呢?nn.Linear()不是也能训练出一个weight吗?

答案是可以的,当然可以直接使用nn.Linear(),只不过输入要改为one-hot Encoding,而不能像nn.Embedding()这种方式直接传入一个index。还有就是需要设置bias=False,因为我们只需要训练一个权重矩阵,不训练偏置

这里给出一个使用单层神经网络来训练Word2Vec的博客

Last Modified: August 2, 2021
Archives Tip
QR Code for this page
Tipping QR Code
Leave a Comment

14 Comments
  1. Young Young

    请问按照文中的代码,负采样可能会取到背景词,这样没关系吗?

    1. mathor mathor

      @Young已修复,感谢~@(呵呵)

  2. w w

    您好,我问一下关于负采样,采样到中心词是没有关系的吗
    还有就是关于取背景词,如果中心词在句子头部或尾部,这样的窗口大小是不够的,按文中的代码来说是取余处理,如果补0(unk代表的idx)的话是否也可以呢

    1. mathor mathor

      @w首先第一个问题,采到中心词是不太好的(因为概率实在太小,所以我代码就没有专门去做判断了)
      第二个问题,补0是可以的

    2. w w

      @mathor十分感谢您的回复@(哈哈)

  3. kyle kyle

    完整的代码链接麻烦修复一下

  4. kou kou

    博主您好,训练完成后除了要保存词向量,是不是还应该保存 word2idx 和 idx2word,否则在另外一个文件加载词向量使用时,找不到输入 word 对应的 idx 和 idx 对应的 word,在新代码里生成 word2idx 和 idx2word 应当是和词向量对不上的吧@(疑问)

    1. mathor mathor

      @kou是的,我这个没考虑那么多

  5. test test

    请问下pos_dot = torch.bmm(pos_embedding, input_embedding)这里对应的意义是什么?
    优化pos_dot和neg_dot对应的实际意义是啥?
    和gensim的预训练相比效率怎么样?@(太开心)
    谢谢~

  6. tzj tzj

    负采样那里我觉得可以优化一下,按照原代码要while循环,很多情况下pos_words(代码里写的pos_indices,是不是写错了?)和neg_words都有<UNK>,导致循环很久才能找出空的交集,可以利用torch.multinomial将中心词和背景词的权重置零,这样负采样就选不到这些词了:
    select_weight = copy.deepcopy(self.word_freq)
    select_weight[pos_words] = 0
    select_weight[center_word] = 0
    neg_words = torch.multinomial(select_weight, self.K * pos_words.shape[0], True)

    1. kk kk

      @tzj直觉上neg_words = torch.multinomial(self.word_freqs, K pos_words.shape[0], True)应该改成neg_indices = torch.multinomial(self.word_freqs, K pos_words.shape[0], True),然后neg_words = self.text_encoded[neg_indices],while循环那里是pos_words代替pos_indices

    2. kk kk

      @kk忽略刚刚的回复,看差了@(汗)

  7. 王志强 王志强

    dataset = WordEmbeddingDataset(text, word2idx, word_freqs)
    dataloader = tud.DataLoader(dataset, batch_size, shuffle=True)
    请问第一行是得到了一个实例,为什么可以通过next(iter())迭代呢?实际上我在我的电脑上并不可以运行代码。谢谢!

    1. 王志强 王志强

      @王志强已经懂了
      整篇博客, 初看不识军真意,在看发现,写的真好!!!
      感谢!