一个基于PyTorch实现的Glove词向量的实例

简介

词向量技术，也称为词嵌入技术(word-embedding)，是一种将高维稀疏的向量压缩到低维稠密向量的技术。常见于自然语言处理领域对单词的预处理过程，例如将单词的one-hot向量是高维稀疏的，不但占用大量空间，而且向量之间提供的信息很少。但经过词嵌入技术生成的向量，不但是低维的并且包含了很多词语间的语义、语法的信息。因此在NLP中，往往会先针对语料生成相应的词向量，然后再把词向量喂入具体任务的神经网络。

Glove词向量是常用的词向量之一，由斯坦福大学发布，原版代码是C实现的，有兴趣的读者可以前往Glove官网下载。本文是根据Glove论文的原理，利用深度学习库PyTorch实现的一个简易版本，仅仅用作学习的示例。

Glove原理

在开始写代码之前，我们对Glove原理进行一次快速介绍。Glove词向量核心思想是利用词与词之间的共现次数来进行训练的，所谓共现就是在一个上下文窗口大小范围内两个词语同时出现的次数，把上下文窗口从语料库的头到尾遍历一次，就得到了一个全局的共现矩阵，然后再基于这个共现矩阵进行词向量的训练。

下面，我们来看看需要优化的目标函数：
$\displaystyle\sum_{i,j}^N f(X_{i,j})(v_i^Tw_j+b_i+b_j-log(X_{i,j}))^2$
其中， $N$ 表示词汇表大小， $X_{i,j}$ 表示词语i和词语j的共现次数， $v_i$ 表示中心词i对应的词向量， $w_j$ 表示上下文词j对应的上下文向量。 $b i 、 b j$ 是偏置项，而 $f (\cdot)$ 表示权重函数，该函数会抑制共现次数过高的词对造成的影响。从目标函数可以看出，模型会训练出两个向量 $v 、 w$ ，原文作者建议将这二者的和作为最后的词向量表示。

实现

1、加载语料库进行预处理
这一步主要是把语料加载到内存中，生成词汇表等。我们直接来看代码：

def getCorpus(filetype, size):if filetype == 'dev':filepath = '../corpus/text8.dev.txt'elif filetype == 'test':filepath = '../corpus/text8.test.txt'else:filepath = '../corpus/text8.train.txt'with open(filepath, "r") as f:text = f.read()text = text.lower().split()text = text[: min(len(text), size)]vocab_dict = dict(Counter(text).most_common(MAX_VOCAB_SIZE - 1))vocab_dict['<unk>'] = len(text) - sum(list(vocab_dict.values()))idx_to_word = list(vocab_dict.keys())word_to_idx = {word:ind for ind, word in enumerate(idx_to_word)}word_counts = np.array(list(vocab_dict.values()), dtype=np.float32)word_freqs = word_counts / sum(word_counts)print("Words list length:{}".format(len(text)))print("Vocab size:{}".format(len(idx_to_word)))return text, idx_to_word, word_to_idx, word_counts, word_freqs

内容很简单，因为语料库也不大，所以能直接加载到内存中。根据语料库内的单词，挑选出前2000个频率最高的词语构成词汇表，剩下的词语则用UNK代替。然后生成一个单词-序号、序号-单词的映射表、词频列表等，以便后续的查询操作。

2、计算词共现矩阵
词共现矩阵的计算方式如下：给出一个句子如"china is a wonderful country and everyone is kind"。假设中心词是wonderful，窗口大小是2，那么我我们增加"wonderful,is"、“wonderful,a”、“wonderful,country”、"wonderful,and"的共现次数，即： $X_{wonderful,is}$ +=1， $X_{wonderful,a}$ +=1、 $X_{wonderful,country}$ +=1等。接着把窗口向右移动，统计下一个中心词-上下文词的共现次数，直到遍历完整个语料库。

def buildCooccuranceMatrix(text, word_to_idx):vocab_size = len(word_to_idx)maxlength = len(text)text_ids = [word_to_idx.get(word, word_to_idx["<unk>"]) for word in text]cooccurance_matrix = np.zeros((vocab_size, vocab_size), dtype=np.int32)print("Co-Matrix consumed mem:%.2fMB" % (sys.getsizeof(cooccurance_matrix)/(1024*1024)))for i, center_word_id in enumerate(text_ids):window_indices = list(range(i - WINDOW_SIZE, i)) + list(range(i + 1, i + WINDOW_SIZE + 1))window_indices = [i % maxlength for i in window_indices]window_word_ids = [text_ids[index] for index in window_indices]for context_word_id in window_word_ids:cooccurance_matrix[center_word_id][context_word_id] += 1if (i+1) % 1000000 == 0:print(">>>>> Process %dth word" % (i+1))print(">>>>> Save co-occurance matrix completed.")return cooccurance_matrix

3、计算权重函数矩阵
由于经常出现的词语实际上不会提供太多的信息，比如to、the之类的停用词，此时他们与别的词的共现频率就会很高，对训练词向量是一个阻碍，因此需要抑制共现频率过高造成的影响，使其在一个合理的范围内。根据原论文所述，其权重函数形式如下：
$\begin{cases} (x/xmax)^{0.75} &\text{if } x < xmax \\ 1 &\text{if } x ≥ xmax \end{cases}$
其中，xmax=100.0

对于 $X_{i,j}$ 可以分别计算出一个对应的 $f(X_{i,j})$ ，由此可以形成一个权重矩阵。实际上权重的计算可以放到前向传播的时候实时计算出来，不需要这样一个矩阵，但笔者这里为了节省前向传播的时间，提前计算出来并保存到矩阵内，等需要用的时候直接查找即可。代码如下所示：

def buildWeightMatrix(co_matrix):xmax = 100.0weight_matrix = np.zeros_like(co_matrix, dtype=np.float32)print("Weight-Matrix consumed mem:%.2fMB" % (sys.getsizeof(weight_matrix) / (1024 * 1024)))for i in range(co_matrix.shape[0]):for j in range(co_matrix.shape[1]):weight_matrix[i][j] = math.pow(co_matrix[i][j] / xmax, 0.75) if co_matrix[i][j] < xmax else 1if (i+1) % 1000 == 0:print(">>>>> Process %dth weight" % (i+1))print(">>>>> Save weight matrix completed.")return weight_matrix

4、创建DataLoader
DataLoader是PyTorch的一个数据加载器，利用它可以很方便地把训练集分批、打乱训练集等，我们新建一个类WordEmbeddingDataset继承自torch.untils.data.Dataset，代码如下图所示：

class WordEmbeddingDataset(tud.Dataset):def __init__(self, co_matrix, weight_matrix):self.co_matrix = co_matrixself.weight_matrix = weight_matrixself.train_set = []for i in range(self.weight_matrix.shape[0]):for j in range(self.weight_matrix.shape[1]):if weight_matrix[i][j] != 0:# 这里对权重进行了筛选，去掉权重为0的项# 因为共现次数为0会导致log(X)变成nanself.train_set.append((i, j))   def __len__(self):'''必须重写的方法:return: 返回训练集的大小'''return len(self.train_set)def __getitem__(self, index):'''必须重写的方法:param index:样本索引 :return: 返回一个样本'''(i, j) = self.train_set[index]return i, j, torch.tensor(self.co_matrix[i][j], dtype=torch.float), self.weight_matrix[i][j]

可以看出，返回的样本包括了词语i和j在词汇表中的需要，以及他们的共现频率和权重，这样就可以直接用于前向传播的计算中。

5、创建训练模型
下面来创建我们的训练模型，在pytorch中创建模型很简单，只需要继承nn.Module，然后重写前向传播函数即可，由于pytorch有自动求导机制，所以我们不需要担心反向传播的问题。

class GloveModelForBGD(nn.Module):def __init__(self, vocab_size, embed_size):super().__init__()self.vocab_size = vocab_sizeself.embed_size = embed_size#声明v和w为Embedding向量self.v = nn.Embedding(vocab_size, embed_size)self.w = nn.Embedding(vocab_size, embed_size)self.biasv = nn.Embedding(vocab_size, 1)self.biasw = nn.Embedding(vocab_size, 1)#随机初始化参数initrange = 0.5 / self.embed_sizeself.v.weight.data.uniform_(-initrange, initrange)self.w.weight.data.uniform_(-initrange, initrange)def forward(self, i, j, co_occur, weight):#根据目标函数计算Loss值vi = self.v(i) #分别根据索引i和j取出对应的词向量和偏差值wj = self.w(j)bi = self.biasv(i)bj = self.biasw(j)similarity = torch.mul(vi, wj)similarity = torch.sum(similarity, dim=1)loss = similarity + bi + bj - torch.log(co_occur)loss = 0.5 * weight * loss * lossreturn loss.sum().mean()def gloveMatrix(self):'''获得词向量，这里把两个向量相加作为最后的词向量:return: '''return self.v.weight.data.numpy() + self.w.weight.data.numpy()

6、训练
下面是对上文的几个点进行串联，最后进行训练，具体的可以参考代码：

EMBEDDING_SIZE = 50      #50个特征
MAX_VOCAB_SIZE = 2000 #词汇表大小为2000个词语
WINDOW_SIZE = 5           #窗口大小为5NUM_EPOCHS = 10         #迭代10次
BATCH_SIZE = 10           #一批有10个样本
LEARNING_RATE = 0.05  #初始学习率
TEXT_SIZE = 20000000  #控制从语料库读取语料的规模text, idx_to_word, word_to_idx, word_counts, word_freqs = getCorpus('train', size=TEXT_SIZE)    #加载语料及预处理
co_matrix = buildCooccuranceMatrix(text, word_to_idx)    #构建共现矩阵
weight_matrix = buildWeightMatrix(co_matrix)             #构建权重矩阵
dataset = WordEmbeddingDataset(co_matrix, weight_matrix) #创建dataset
dataloader = tud.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
model = GloveModelForBGD(MAX_VOCAB_SIZE, EMBEDDING_SIZE) #创建模型
optimizer = torch.optim.Adagrad(model.parameters(), lr=LEARNING_RATE) #选择Adagrad优化器print_every = 10000
save_every = 50000
epochs = NUM_EPOCHS
iters_per_epoch = int(dataset.__len__() / BATCH_SIZE)
total_iterations = iters_per_epoch * epochs
print("Iterations: %d per one epoch, Total iterations: %d " % (iters_per_epoch, total_iterations))
start = time.time()
for epoch in range(epochs):loss_print_avg = 0iteration = iters_per_epoch * epochfor i, j, co_occur, weight in dataloader:iteration += 1optimizer.zero_grad()   #每一批样本训练前重置缓存的梯度loss = model(i, j, co_occur, weight)    #前向传播loss.backward()     #反向传播optimizer.step()    #更新梯度loss_print_avg += loss.item()
torch.save(model.state_dict(), WEIGHT_FILE)

实验结果

1、利用余弦相似度衡量向量之间的相似性
按照上述代码进行训练后，我们执行以下代码来看看训练得到的模型的效果怎么样。我们利用余弦相似度来衡量两个向量是否相似，如果两个词语是同属一类的词语，那么它们的向量的距离应该很近的。我们定义以下的函数：

def find_nearest(word, embedding_weights):index = word_to_idx[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]]#找到前10个最相近词语

然后列举几个词语，看看跟他们最相似的几个词语是什么，运行以下代码：

glove_matrix = model.gloveMatrix()
for word in ["good", "one", "green", "like", "america", "queen", "better", "paris", "work", "computer", "language"]:print(word, find_nearest(word, glove_matrix))

运行结果如下：

可以看出，在某些词语比如数字、颜色、计算机、形容词比较级等，取得了不错的效果。

2、利用向量之间的运算考察词向量的关联性
词向量有另外一个特性，即“man to king is woman to ?”，向量化表述就是 $v (c a n d i d a t e) = v (k i n g) - v (m a n) + v (w o m a n)$ ，也即是候选词的词向量可以由另外三个词向量通过加减运算得到。那么，显然上述的候选词应该是“queen”。我们利用这个特性来举几个例子：

def findRelationshipVector(word1, word2, word3):word1_idx = word_to_idx[word1]word2_idx = word_to_idx[word2]word3_idx = word_to_idx[word3]embedding = glove_matrix[word2_idx] - glove_matrix[word1_idx] + glove_matrix[word3_idx]cos_dis = np.array([scipy.spatial.distance.cosine(e, embedding) for e in glove_matrix])for i in cos_dis.argsort()[:5]:print("{} to {} as {} to {}".format(word1, word2, word3, idx_to_word[i]))findRelationshipVector('man', 'king', 'woman')
findRelationshipVector('america', 'washington', 'france')
findRelationshipVector('good', 'better', 'bad')

运行以上代码，可以得到以下结果：

可以看出，结果一般般，但还是能看得到模型在努力地筛选出了更为相关的信息，只不过准确度需要再进一步提高。

3、利用SVD进行词向量的降维并可视化显示
SVD技术常用于将高维向量降维，我们将词向量降到二维，然后把降维后的向量画到画板上，看看这些词向量有没有很好地聚合在一起。如果聚合在一起了，表示词向量是相近的，也即是模型有比较好地学习到了词语之间的语义、语法关系。运行以下代码并观察结果：

#数据降维以及可视化
candidate_words = ['one','two','three','four','five','six','seven','eight','night','ten','color','green','blue','red','black',                      'man','woman','king','queen','wife','son','daughter','brown','zero','computer','hardware','software','system','program','america','china','france','washington','good','better','bad']
candidate_indexes = [word_to_idx[word] for word in candidate_words]
choosen_indexes = candidate_indexes
choosen_vectors = [glove_matrix[index] for index in choosen_indexes]U, S, VH = np.linalg.svd(choosen_vectors, full_matrices=False)
for i in range(len(choosen_indexes)):
plt.text(U[i, 0], U[i, 1], idx_to_word[choosen_indexes[i]])coordinate = U[:, 0:2]
plt.xlim((np.min(coordinate[:, 0]) - 0.1, np.max(coordinate[:, 0]) + 0.1))
plt.ylim((np.min(coordinate[:, 1]) - 0.1, np.max(coordinate[:, 1]) + 0.1))
plt.show()

从结果图可以观察到，同类型语义相近的词语基本都靠得很近。

结论

通过对实验结果的观察，可以发现训练得到的Glove模型效果还不错，由于笔者设备的限制，仅仅对1500万个词语的语料库以及2000个词语的词汇表，采取了50个特征值，进行了10个轮次的训练。如果语料库更大、词汇表更大以及特征数量更多，相信训练出来的词向量将会有更好的效果。由于本文仅仅是作为一个学习的样例，所以Glove的训练到此为止，谢谢各位的阅读~

项目代码

最后贴出笔者的项目代码地址，包括了语料库以及模型等，有需要的同学可以自行下载：PyTorchGlove