本文主要是使用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)如果使用与训练好的词向量，则采用
代码语言：javascript复制pretrained_weight = np.array(pretrained_weight)
embeds.weight.data.copy_(torch.from_numpy(pretrained_weight))想要查看某个词的词向量，需要传入这个词在词典中的index，并且这个index得是LongTensor型的
代码语言：javascript复制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, embedding_dim)
上图的流程是把文章中的单词使用词向量来表示
- 提取文章所有的单词，把所有的单词按照频次降序排序（取前4999个，表示常出现的单词。其余所有单词均用'<UNK>'表示。所以一共有5000个单词）
- 500个单词使用one-hot编码
- 通过训练会生成一个$5000times 300$的矩阵，每一行向量表示一个词的词向量。这里的300是人为指定，想要每个词最终编码为词向量的维度，你也可以设置成别的
这个矩阵如何获得呢？在Skip-gram模型中，首先会随机初始化这个矩阵，然后通过一层神经网络来训练。最终这个一层神经网络的所有权重，就是要求的词向量的矩阵
从上面的图中看到，我们所学习的embedding层是一个训练任务的一小部分，根据任务目标反向传播，学习到embedding层里的权重weight。这个weight是类似一种字典的存在，他能根据你输入的one-hot向量查到相应的Embedding vector
Pytorch实现
导包
代码语言：javascript复制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 math
import pandas as pd
import scipy
import sklearn
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
文件中的内容是英文文本，去除了标点符号，每个单词之间用空格隔开
读取文本数据并处理
代码语言：javascript复制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>"
idx2word = [word for word in vocab_dict.keys()]
word2idx = {word:i for i, word in enumerate(idx2word)}
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需要以下内容：
- 把所有word编码成数字
- 保存vocabulary，单词count、normalized word frequency
- 每个iteration sample一个中心词
- 根据当前的中心词返回context单词
- 根据中心词sample一些negative单词
- 返回单词的counts
为了使用DataLoader，我们需要定义以下两个function
__len__()：返回整个数据集有多少item__getitem__(idx)：根据给定的idx返回一个item
这里有一个好的tutorial介绍如何使用PyTorch DataLoader
代码语言：javascript复制class WordEmbeddingDataset(tud.Dataset):
def __init__(self, text, word2idx, idx2word, word_freqs, word_counts):
''' text: a list of words, all text from the training dataset
word2idx: the dictionary from word to index
idx2word: index to word mapping
word_freqs: the frequency of each word
word_counts: the word counts
'''
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.idx2word = idx2word
self.word_freqs = torch.Tensor(word_freqs)
self.word_counts = torch.Tensor(word_counts)
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]是正确单词数量
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
代码语言：javascript复制dataset = WordEmbeddingDataset(text, word2idx, idx2word, word_freqs, word_counts)
dataloader = tud.DataLoader(dataset, batch_size, shuffle=True)可以随便打印一下看看
代码语言：javascript复制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模型
代码语言：javascript复制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.in_embed(pos_labels)# [batch_size, (window * 2), embed_size]
neg_embedding = self.in_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.numpy()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])训练模型
代码语言：javascript复制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_embeddings()
torch.save(model.state_dict(), "embedding-{}.th".format(EMBEDDING_SIZE))如果没有GPU，训练时间可能比较长
词向量应用
我们可以写个函数，找出与某个词相近的一些词，比方说输入good，他能帮我找出nice，better，best之类的
代码语言：javascript复制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 [idx_to_word[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']
