基于PyTorch实现word2vec模型

word2vec是google于2013年的《Distributed Representations of Words and Phrases and their Compositionality 》以及后续的《Efficient Estimation of Word Representations in Vector Space 》两篇文章中提出的一种高效训练词向量的模型, 基本出发点是上下文相似的两个词,它们的词向量也应该相似, 比如香蕉和梨在句子中可能经常出现在相同的上下文中,因此这两个词的表示向量应该就比较相似。本文前部分主要从理论方面介绍word2vec,后半部分主要基于PyTorch框架实现word2vec模型(skip-gram)。

word2vec理论

word2vec模型中比较重要的概念是词汇的上下文, 说白了就是一个词周围的词, 比如$w_t$的范围为1的上下文就是$w_{t-1}$和$w_{t+1}$,在word2vec中提出两个模型(假设上下文窗口为3)

  • CBOW(Continuous Bag-of-Word): 以上下文词汇预测当前词: $w_{t-1}$和$w_{t+1}$去预测$w_t$
  • SkipGram: 以当前词预测其上下文词汇: $w_t$去预测$w_{t-1}$和$w_{t+1}$。

两个模型结构如下图所示:

下面将会从最简单的上下文只有一个词的情形入手, 然后扩展到CBOW以及Skip-gram, 介绍原理以及参数训练过程. 关于word2vec的训练这里将会从完全的BP神经网络的过程来介绍.
One-Word Model

One-Word Model

首先先看简化版入手: 输入输出都只有一个词, 如下图示:

首先说明符号:

  • V: 词汇表长度
  • N: 隐层神经元个数, 同时也是词向量维度
  • $W \in \mathcal{R}^{V\times N}$: 输入层到隐层的权重矩阵, 其实就是词向量矩阵,其中每一行代表一个词的词向量
  • $W^{‘} \in \mathcal{R}^{N\times V}$ 隐层到输出层的权重矩阵, 其中每一列也可以看作额外的一种词向量

下面从神经网络的前向过程开始介绍:

我们需要做的是用输入的词去预测输出的词. 其中 输入层的单词$w_I$使用one-hot来表示的, 即在上图中$x_1, x_2, x_3,…,x_V$只有$x_k$为1, 其余为0, 其中k可以是输入的词在词汇表中的索引下标. 之后就是经过词向量矩阵W 连接输入层和隐层. 其中由于X中只有一个1, 因此经过与W相乘, 相当于取出W中的的第k行,实际也就是输入单词的$w_I$的N维的词向量,使用$v_{w_I}$表示,来作为隐层的值,注意word2vec的隐层并没有激活函数:

然后考虑从隐层的h到输出层Y, 同样h经过矩阵$W^{‘}$相乘,得到一个V×1的向量u:

其中u每个元素$u_j$ 就是$W^{‘}$的第j列用$v_{w_j}^{‘}$表示, 与h做内积得到: $u_j = v_{w_O}^{‘T}\cdot h$,含义就是词汇表中第j个词的分数,我们的目的就是要根据输入词$w_I$去预测输出的词,因此预测的词就取分数最高的即可,这里为了方便概率表示,使用softmax将u归一化到[0,1]之间, 从而作为输出词的概率, 其实是一个多项分布, 也就是上图中的y:

其中$v_w$与$v_{w}^{′}$ 都称为词w的词向量,一般使用前者作为词向量,而非后者,原因后续会解释。至此前向过程完成,就是给定一个词作为输入,来预测它的上下文词,还是比较简单的,属于简化版的神经语言模型。这个过程中需要用到的参数有两个词向量矩阵W,$W^{‘}$,下面就是重点了,介绍如何根据语料库来训练模型,更新参数,得到最终的词向量。

训练

首先明确训练数据的格式,对于一个训练样本$(w_i,w_o)$,输入是词$w_i$的one-hot 的维度为V的向量x,模型预测的输出同样也是一个维度为V的向量y, 同时真实值$w_o$也是用one-hot表示,记为$t=[0,0,0,…1,0,0]$,其中假设$t_{j^{\star}}=1$, 也就是说$j^{\star}$是真实单词在词汇表中的下标,那么根据最大似然或者上面的语言模型,目标函数可以定义如下:

一般我们习惯于最小化损失函数,因此定义损失函数:

然后结合反向传播一层层求梯度,使用梯度下降来更新参数。

先求隐层到输出层的向量矩阵$W^{‘}$的梯度:

这里面的$y_j$和$t_j$分别是预测和真实值的第j项,$h_i$是隐层的第i项。

考虑$\frac{\partial E}{\partial u_j} = y_j - t_j$: 直接对原始求导,如下:

先考虑E的对数部分: $\frac{\partial \log \sum\exp(u_k)}{\partial u_j} = \frac{\exp(u_j)}{\sum \exp(u_k)} = y_j$

再看 $ -u_{j^{\star}}$对$u_j$的梯度:当$j = j^{\star}$时,此时梯度为$=-1 = -t_j$,反之二者不等时,

此时梯度为$=0= -t_j$,所以综合求导$\frac{\partial E}{\partial u_j} = y_j - t_j$,这个减法可以理解为输出层的第j项为预测值与真实值的差

因此梯度下降更新公式为:

整合为$W^{‘}$的列向量$\mathbf{v^{‘}}_{w_j} = \{w^{‘}_{ij}| i=1,2,3,…,N\}$的形式如下:

也就是说对每个训练样本都需要做一次复杂度为V的操作去更新$W^{‘}$

接着考虑隐层h的更新,其实也是输入层到隐层的矩阵W的更新,继续反向传播,跟神经网络的相同, 输出层的V个神经元都会影响$h_i$:

其中$W_i^{‘}$是$W^{‘}$的第i行, 这里为了方便书写, 令$P = \{ y_j - t_j | j= 1,2,3,..,V\}$, 因此整合成整个隐层的向量h:

得到一个N维的向量,上面已经介绍过,h就是词向量矩阵W的一行: $ = W^T X = v_{w_i}^T$ ,但是因为X中只有一个1,因此每次只能更新的一行$v_{w_i}$,其余行的梯度 = 0,所以$v_{w_i}$的更新公式为:

到此为止, 一个训练样本的反向传播训练过程就为止了。 我们可以看到,对于输入层到隐层的矩阵W,我们每次训练只需要更新一行向量即可,而对于隐层到输出层的矩阵$W^{‘}$的所有N×V个元素都需要更新一遍,这里的计算量还是很大的。

这一节主要比较细致的介绍了最简单的输入输出只有一个单词的情况的推理和训练的过程,后面的CBOW(上下文预测单词)以及SG(单词预测上下文)均基于这一节扩展开来。

CBOW Model

这一部分讲word2vec的第一个形式: Continurous Bag-Of-Word,模型图示如下:

跟上一个模型唯一的不同就是输入不再是一个词$w_i$, 而是多个词,上图中一共有C个单词: $x_{1k},x_{2k},…,x_{Ck}$,每个x都是one-hot表示。 这样隐层的h的计算就会不同了,之前一个单词的模型是直接取出W的一行$v_{w_i}$作为h的值,在CBOW中则是取出W中输入的所有C个单词的词向量,然后直接取平均,如下:

后面隐层到输出层的过程与One-Word Model 一模一样,包括目标函数定义, 反向传播训练等。将$W^{‘}$的更新公式照抄下来如下,依旧是每次都需要更新所有的行:

隐层神经元的梯度也相同:

下面考虑输入层到隐层稍微有些不同,在One-Word Model里面因为输入只有一个词,因此每次训练只更新这个词对应到W的那一行,但是在CBOW里面有多个词,这里采取的策略是将h的梯度均摊到每个词上,因此每次训练会更新W中的C行,如下:

到此为止 CBOW 的推理和训练过程也介绍完毕,基本跟One-Word Model 一样。

SkipGram Model

现在开始介绍word2vec的第二种形式: SkipGram(根据单词预测上下文),这个模型与One-Word Model不同的地方在于,SG的输出有多个词,而非One-Word 中输出只有一个词,这样输出层就不是一个多项分布了,而是C个多项分布了,模型图示如下:

因此从输入层到隐层部分与One-Word Model 相同,隐层神经元的计算方式如下:

因为输出层是有C个单词, 因此有C个多项分布: $y_{1}, y_{2}…y_{C}$, 因此前向计算的过程也需要分开计算,如下公式,用来计算第c输出单词的预测的多项分布中第j项,相比One-Word Model 多了一个c参数:

需要主要的是这C个输出向量是相互独立的,可以当做是独立的C个One-Word Model 中的输出向量,相互之间没有影响,并且从图中也可以看出,连接隐层与C个输出层的参数矩阵$W^{‘}$是共享的,于是便有: $u_{c,j} = u_{j} ={v^{‘}_{w_j}}^T \cdot \mathbf{h}$ 这里的$v_{w_j}^{‘}$的含义与One Word Model 中相同,都代表$W^{‘}$的第j列,同时也是词汇表中第j个单词的一种词向量(虽然实际中不用)。

从前向后 根据上述公式计算出C个输出向量之后,在每个V维向量中选取概率最大的作为输出的单词,这样根据输出单词wI 就得到了C个输出单词,也就达到了根据单词预测上下文的目的。

下面开始介绍SG的反向传播训练的过程,这个跟前面的有些许的不同, 首先是损失函数:

前面说过输出的C个词是相互独立,因此$P(w_1,w_2,…,w_C|W_i) = \Pi P(w_c|w_i)$, 此外$j_c^{\star}$的含义同One-Word Model 中的$u_j^{\star}$一样,都代表训练的真实的输出单词在词汇表的下标。下面开始从后向前求梯度,对第c个词对应的多项分布的第j项的梯度:

然后考虑$W^{‘}$的梯度,考虑到C个多项分布产生影响,因此需要求和:

跟CBOW一样,为了方便书写定义$Q_j = \sum_{c=1}^{C}(y_{c,j} - t_{c,j}), j = 1,2,3,…,V$,矩阵Q的维度是V1有了梯度,就可以利用梯度下降 更新$w_{ij}^{‘}$:

或者写成词向量的形式,其实就是$W^{‘}$的一列:

接着考虑对隐层神经元的梯度:

因此跟One-Word Model一样整合成向量的形式:

考虑到输入只有一个词,因此跟One-Word Model 一样:$\mathbf{h} = {v_{w_I}}^T$, 因此每次训练更新词向量矩阵W的一行:

到此SkipGram模型的前向推理与后向参数的训练也介绍完毕了。

优化

复杂度

前面的CBOW与SG模型是标准的word2vec模型,或者说是神经语言模型的简化版,去掉了隐层的激活函数,其余的变化不大,因此训练效率还是很低的。我们分析下训练的复杂度。首先明确需要学习的两个词向量矩阵W,$W^{‘}$,从前面的推导中知道对于每一个训练样本,CBOW更新W的C行,SG更新W其中一行,也就是每次更新有限个词的词向量。但是对于$W^{‘}$则不同了,正如前面一直提到的,无论是CBOW还是SG,对每个训练样本(或者Mini Batch)从梯度更新中需要对$W^{‘}$的所有V×N个元素,也就是词汇表中所有V个单词都需要更新词向量,考虑现实任务词汇表一般是几十万,上百万千万级别的, 这个计算成本是巨大的。

关于计算成本大的原因,除了上面提到的训练部分,还有就是在每次前向计算的时候,隐层到输出层的softmax函数计算输出层V个元素,计算量也是很大,这样整个模型现实意义不大。

考虑到计算量大的部分都是在隐层到输出层上,尤其是$W^{‘}$的更新。因此word2vec使用了两种优化策略:

  • Hierarchical Softmax
  • Negative Sampling。

二者的出发点一致,就是在每个训练样本中,不再完全计算或者更新$W^{‘}$这个矩阵。二者都不再显示使用$W^{‘}$这个矩阵。因此这也就解释了前面说的为什么不用$W^{‘}$作为最终词向量。

在多一句,其实上述训练和推理的复杂度很大的根本原因是softmax的分母上的$\sum$,因此在求梯度的时候,就会有V次的计算。因此下面的两种方法其实是对softmax的优化,不仅仅限制在word2vec.

两种优化方式使得word2vec的训练速度大大提升,并且词向量的质量几乎没有下降,这也是word2vec在NLP领域如此流行的原因。 下面依次介绍这两种优化算法。

Hierarchical SoftMax

首先Hierarchical SoftMax(HS)并不是word2vec提出来的, 而是之前Bengio在2005年最早提出来专门为了加速计算神经语言模型中的softmax的一种方式, 这里介绍如何在word2vec中使用. HS主要基于哈夫曼树(一种二叉数)将计算量大的部分变为了一种二分类的问题. 先看下面的图, 原来的模型在隐层之后通过$W^{‘}$连接输出层, 现在HS则去掉了$W^{‘}$, 隐层h直接与下面的二叉树的root节点相连:

其中图中白色的叶子节点表示词汇表中所有的|V|个词, 黑色节点表示非叶子节点, 每一个叶子节点也就是每一个单词, 都对应唯一的一条从root节点出发的路径. 而我们的目的是使的$w=w_o$这条路径的概率最大,即: $P(w=w_o|w_i)$ 最大, 此时每一个分支都代表一个选择, 向左转还是向右转. 所以如何判断向左还是向右呢? 我们用n(w,j)表示从root到叶子节点w的路径上的第j个非叶子节点, 并且每个非叶子节点都对应一个向量$v_{n(w,j)^{‘}}^{‘}$, 维度与h相同, 然后使用一个sigmod函数: $\sigma(x)= \frac{1}{1+\exp(-x)} \in [0, 1]$, 结合向量的内积, 来判断该向左还是向右, 如下, 第n个节点向左 以及向右的概率定义:

有了上述的概率, 我们可以重新定义$P(w_o|w_i)$了:

其中I()是指示函数, 条件成立值为1, 反之为-1. 而L(w)表示整条路径的长度, 这样整个概率就是从root节点到叶子节点这条路径的概率, 这样我们在训练的时候, 通过训练样本来更新非叶子节点的参数v′w.

举个例子, 比如上图中的加粗的黑色路径: $(n(w_2,1), n(w_2, 2), n(w_2, 3), w_2$, 就是说假设有一个训练样本是$(w_i,w2)$, 我们需要使得$P(w_o=w_2|w_i)$概率最大:

并且在一个非叶子节点处, 向左向右的概率和为1, 因此一直分裂下去,最后的和肯定还是1. 因此可以很容易得到:

这一点的证明是有必要的, 因为在原始的softmax本身保证了所有单词的概率和是1, 而通过上式也知道了通过HS得到的输出层仍然是一个概率多项分布, 输出所有的单词概率和为1.

讲完了从前向后的如何得到输出单词的概率的过程, 下面开始后向传播的训练过程. 首先需要明确的是训练的参数: 输入层与隐层的词向量矩阵W, 以及二叉树的非叶子节点对应的向量$\{v{‘}_{n(w,j)}, j =1,2,3,..,L(w)-1\}$.

为了书写方便,下面简化一部分的符号: 用[I]表示前面的指示函数$I(n(w,j+1)==left)$, 使用v′j表示$v^{‘}_{n(w,j)}$

对于一组训练数据, 损失函数的定义与前面相同, 最大似然(注意这里以One-Word Model为例,CBOW与Skip-Gram按照上述的思路完全一样):

之后便可以逐项求梯度了, 先考虑$v^T \mathbf{h}$, 注意$\frac{\partial \sigma(x)}{x} = \sigma(1-\sigma)$

之后对[I]分情况讨论, 可得: $\frac{\partial E}{\partial {v^{‘}_{j}}^T \mathbf{h}} = \sigma( {v^{‘}_{j}}^T \mathbf{h}) - t_j$如果[I]=1, 那么$t_j=1$, 否则$ tj=0$ , 这个公式与前面的$y_j−t_j$很类似, 可以理解为预测值与实际值的差别。

有了上述的梯度,就可以很简单的求出v′的梯度了:

有了梯度,便可以更新了, 具体公式还是梯度下降:

也就是说对于一个训练样本, 我们只需要更新L(w)−1个向量就好了, 而未优化的版本需要更新V个, 相当于时间复杂度从O(V)降到了O(logV) , 这个提升还是非常大的, 同时在考察空间复杂度,HS的二叉树的非叶子节点有V−1个,也就是我们需要V−1存储$v_j^{‘}, j=1,2,3..V−1$ ,优化之前则是V个, 空间复杂度相同, 但是时间复杂度却大大降低。

然后考虑隐层h的梯度,因为我们的优化目标都是在隐层到输出层,因此前面的几乎不变, 跟One-Word Model 一样,路径上的非叶子节点的表达式都含有h,因此需要对梯度求和:

其实跟前面很一样了, 只需要替代下导数就好了, 后面就不再赘述了。

整个Hierarchical Softmax的优化算法介绍完了,隐层到输出层的计算量从O(V), 利用二叉树(更加具体来说是哈夫曼树)降为了O(logV)

Negative Sampling

相比Hierarchical Softmax用复杂的树形结构来对softmax进行优化, 负采样(Negative Sampling)更加直接简单。因为softmax的存在,在使用梯度下降的时候,每个训练样本需要更新V个向量,这是根本问题。因此负采样NS仅仅更新一小部分向量,一般是5-20个,可以手工设定,而非全部V个(一般来说V都是几万到几百万的量级)。

再来考虑原始模型的训练过程,对于一个训练样本$(w_i,w_o)$, 我们要使得P(w=wO|wI)最大, 也就是说要使得输出层的其它单词的概率要比较小一些,基于这种思想, 我们将输出层的V个样本分为正例(Positive Sample)也就是$w_o$对应的项, 以及剩余V−1个负例(Negative Samples)。举个例子有个样本phone number,这样$w_i$=phone,$w_o$=number, 正例就是number这个词, 负例就是不太可能与phone共同出现的词。

负采样的思想便是每次训练只随机取一小部分的负例使他们的概率最小,以及对应的正例概率最大。那么如何对负例进行抽样呢? 这里就需要定义一个noise distribution,有了分布Pn(w)就可以依据概率分布进行带权随机抽样了。 在word2vec中,作者直接使基于词的频次的词的权重分布:

相比于直接使用频次作为权重, 取0.75幂的好处可以减弱不同频次差异过大带来的影响,使得小频次的单词被采样的概率变大。

下面基于上面的思想,直接给出具体的损失函数:

Wneg是负例单词集合。 上述公式跟原本的负采样的有些许差别,具体的推导细节可以参考这篇Goldberg, Y. and Levy, O. (2014). word2vec explained: deriving mikolov et al.’s negative- sampling word-embedding method. 这里不再赘述了。 只需要记住,NS也是对softmax函数做的优化,除了word2vec,在其他地方涉及到softmax的均可以采用类似的思想来重写目标函数。

有了损失函数,然后跟HS一样,用梯度下降的方法更新v′即可,推导思路跟HS也是一样,分$w_j$是否等于$w_o$来讨论,最终得到的梯度公式与HS一模一样:

这里的tj含义跟前面提到的几次都相同。最后给出更新v′的公式:

从公式也可以看出,对于每个训练数据,每次只需要更新很少的几个向量,而非原始的V

个,大大降低了训练复杂度。

跟HS相同,NS的优化也是对隐层到输出层的优化,因此前面输入层到隐层的梯度只需要做相应修改即可,没有影响,具体可以参考HS的公式,这里不再重复了。

到此为止,两种针对softmax的优化方式介绍完毕,简单对比一下两种方式:

  • S实现比较复杂, 需要借助哈夫曼树,而且HS是专门优化softmax的方法
  • S相对简单,仅仅是简单采样,NS则是一种概率采样的方法,属于NCE(Noise Contrastive Estimation)的一种简单形式,代码实现也很简单。 训练得到的词向量质量也很高,相对常用一些。

PyTorch实现

前部分,我们详细地介绍了Word2vec的理论知识,接下来,我们将基于PyTorch框架实现word2vec模型(skip-gram).

我们按照如下目录结构进行整个项目的开发和实验:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
├── pyword2vec
| └── callback
| | └── lrscheduler.py  
| └── config
| | └── word2vec_config.py  
| └── dataset           
| └── io              
| └── model
| └── output           
| └── preprocessing    
| └── train
| └── utils
├── get_similar_words.py
├── train_gensim_word2vec.py
├── train_word2vec.py

备注:代码中也提供了gensim版本训练word2vec,感兴趣的可以拿两个版本结果进行比较,当然我们实现的PyTorch版本可能从效果上会低于gensim版本,因为gensim版本做了很多优化。

前面提到,skipgram模型是使用当前词预测其上下文,那么我们将按照下面格式进行构造训练数据集:

备注: 没有把全部完整的代码贴出,只是挑选了一些进行说明。完整代码地址可从github获取。

在本次实验中,我们使用的是公开的zhihu数据集,本身已经分好词,且一行表示一个句子,因此,我们直接加载数据集,如下所示:

1
2
3
4
5
6
7
8
9
10
11
    # 读取数据,并进行预处理
def build_examples(self):
self.examples = []
with open(self.data_path, 'r') as fr:
for i, line in tqdm(enumerate(fr), desc='read data and processing'):
# 数据首行为列名
if i == 0 and self.skip_header:
continue
line = line.strip("\n")
if line:
self.examples.append(self.split_sent(line))

当全部数据集加载到内存中,我们需要构造vocab,保存word与id之间的映射关系,之后在将句子转化为id序列时候会用到,即:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
    # 建立语料库
def build_vocab(self):
count = Counter()
for words in tqdm(self.examples,desc = 'build vocab'):
count.update(words)
count = {k: v for k, v in count.items()}
count = sorted(count.items(), key=operator.itemgetter(1),reverse=True)
all_words = [(w[0],w[1]) for w in count if w[1] >= self.min_freq]
if self.vocab_size:
all_words = all_words[:self.vocab_size]
all_words = all_words+[('<unk>',0)]
word2id = {k: (i,v) for i,(k, v) in zip(range(0, len(all_words)),all_words)}
self.word_frequency = {tu[0]: tu[1] for word, tu in word2id.items()}
self.vocab = {word: tu[0] for word, tu in word2id.items()}
pkl_write(data = word2id,filename=self.vocab_path)

在实际实现word2vec模型时,做了很多细节的优化,这里我们首先构建负样本抽样分布,即:

1
2
3
4
5
6
7
8
9
10
11
    # 构建负样本
def build_negative_sample_table(self):
self.negative_sample_table = []
sample_table_size = 1e8
pow_frequency = np.array(list(self.word_frequency.values())) ** 0.75
words_pow = sum(pow_frequency)
ratio = pow_frequency / words_pow
count = np.round(ratio * sample_table_size)
for wid, c in enumerate(count):
self.negative_sample_table += [wid] * int(c)
self.negative_sample_table = np.array(self.negative_sample_table)

接着对全量数据集进行采样,降低高频词的出现,即:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
    # 数据采样,降低高频词的出现
def subsampling(self,total = 2 ** 32):
pow_frequency = np.array(list(self.word_frequency.values()))
words_pow = sum(pow_frequency)
ratio = pow_frequency / words_pow
delete_int = [self.reserve_ratio(p,total = total) for p in ratio]

self.train_examples = []
for example in self.examples:
words = [self.vocab[word] for word in example if
word in self.vocab and delete_int[self.vocab[word]] >= random.random() * total]
if len(words) > 0:
self.train_examples.append(words)
del self.examples

当全部数据操作处理完之后,我们将数据集按照skip-gram模型要求格式进行构造训练样本,即:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
    # 构建skip gram模型样本
def make_iter(self):
for example in self.train_examples:
if len(example) < 2:
continue
reduced_window = self.random_s.randint(self.window_size)
for i,w in enumerate(example):
words_num = len(example)
window_start = max(0, i - self.window_size + reduced_window)
window_end = min(words_num, i + self.window_size + 1 - reduced_window)
pos_v = [example[j] for j in range(window_start, window_end) if j != i]
pos_u = [w] * len(pos_v)
neg_u = [c for c in pos_v for _ in range(self.negative_num)]
neg_v = [v for u in pos_u for v in self.get_neg_word(u)]
yield pos_u,pos_v,neg_u,neg_v

以上我们简单地完成了整个skipgram模型的训练样本构造。接下来,我们定义整个skipgram模型,即:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
#encoding:Utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F

class SkipGram(torch.nn.Module):
def __init__(self, embedding_dim, vocab_size):
super(SkipGram, self).__init__()
initrange = 0.5 / embedding_dim
self.u_embedding_matrix = nn.Embedding(vocab_size,embedding_dim)
self.u_embedding_matrix.weight.data.uniform_(-initrange,initrange)
self.v_embedding_matrix = nn.Embedding(vocab_size,embedding_dim)
self.v_embedding_matrix.weight.data.uniform_(-0, 0)

def forward(self, pos_u, pos_v,neg_u, neg_v):
embed_pos_u = self.v_embedding_matrix(pos_u)
embed_pos_v = self.u_embedding_matrix(pos_v)
score = torch.mul(embed_pos_u, embed_pos_v)
score = torch.sum(score,dim = 1)
log_target = F.logsigmoid(score).squeeze()

embed_neg_u = self.u_embedding_matrix(neg_u)
embed_neg_v = self.v_embedding_matrix(neg_v)

neg_score = torch.mul(embed_neg_u,embed_neg_v)
neg_score = torch.sum(neg_score, dim=1)
sum_log_sampled = F.logsigmoid(-1 * neg_score).squeeze()

loss = log_target.sum() + sum_log_sampled.sum()
loss = -1 * loss
return loss

最后,在整个项目根目录中,运行以下命令,完成整个PyTorch版本的word2vec模型训练:

1
python train_word2vec.py

结果

大概 6 次 epochs 之后,可得到以下结果:

目标词Top10目标词Top10
中国中国 : 1.000男人男人 : 1.000
中国美国 : 0.651男人女人 : 0.764
中国日本 : 0.578男人女生 : 0.687
中国国家 : 0.560男人男生 : 0.670
中国发展 : 0.550男人喜欢 : 0.625
中国文化 : 0.529男人恋爱 : 0.601
中国朝鲜 : 0.512男人岁 : 0.590
中国经济 : 0.504男人女 : 0.588
中国世界 : 0.493男人感觉 : 0.586
中国社会 : 0.481男人男朋友 : 0.581

从结果中看,整体上效果还不错。

完成代码地址:github

感兴趣的小伙伴可以下载下来跑跑或者学习。

参考内容

-------------本文结束感谢您的阅读-------------
;