引言
本着“凡我不能创造的，我就不能理解”的思想，本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架，该框架类似PyTorch能实现自动求导。
要深入理解深度学习，从零开始创建的经验非常重要，从自己可以理解的角度出发，尽量不使用外部完备的框架前提下，实现我们想要的模型。本系列文章的宗旨就是通过这样的过程，让大家切实掌握深度学习底层实现，而不是仅做一个调包侠。
前面介绍的CBOW和Skip-gram模型有一个重大的缺点，就是计算量太大了。主要是在最终的多分类问题上，我们经过了一个Softmax操作，想象一下百万级的词汇量，那么Softmax需要计算百万次。
针对这个问题有两种优化方法，分别是层次Softmax和负采样。本文我们主要介绍带负采样的Skip-gram。
负采样
SGNS(Skip-Gram with Negative-Sampling)，即带负采样的Skip-gram。
它将多分类任务简化为二分类任务，即不预测每个单词附近会出现某个单词，而是判断某个单词是否会在 w w w附近出现。同样，训练完之后，我们需要的是学习到的权重。
二分类任务其实就是一个逻辑回归分类器，它的训练过程如下：
- 将目标词和一个上下文单词组成正例
- 随机采样词典中的其他单词与目标词组成负例
- 训练逻辑回归分类器去区分正例和负例
- 使用学到的权重作为嵌入
逻辑回归分类器
假设窗口大小为 2 2 2，目标词为 w t w_t wt?，对于句子 ? ? w t ? 2 ? w t ? 1 ? w t ￣ ? w t + 1 ? w t + 2 ? ? \cdots \, w_{t-2} \, w_{t-1} \, \underline{w_t} \, w_{t+1} \, w_{t+2} \, \cdots ?wt?2?wt?1?wt??wt+1?wt+2??。
它的上下文单词为
c
∈
{
w
t
?
2
,
w
t
?
1
,
w
t
+
1
,
w
t
+
2
}
c \in \{w_{t-2},w_{t-1},w_{t+1},w_{t+2}\}
c∈{wt?2?,wt?1?,wt+1?,wt+2?}，假设
c
c
c代表其中任意上下文单词，和目标词组成元组
(
w
t
,
c
)
(w_t,c)
(wt?,c)。那么分类器输出
c
c
c是
w
t
w_t
wt?上下文单词的概率：
P
(
+
∣
w
t
,
c
)
(1)
P(+|w_t,c) \tag{1}
P(+∣wt?,c)(1)
为了让它是一个概率，那么如果
c
c
c不是上下文的概率就可以用
1
1
1减去上式得到，以保证这两个事件概率之和为
1
1
1：
P
(
?
∣
w
t
,
c
)
=
1
?
P
(
+
∣
w
t
,
c
)
(2)
P(-|w_t,c) = 1 - P(+|w_t,c) \tag {2}
P(?∣wt?,c)=1?P(+∣wt?,c)(2)
现在问题是我们如何计算这个概率呢？可能你已经看出来了，对，就是通过Sigmoid函数。具体做法为，还是计算
w
t
w_t
wt?和
c
c
c这两个词嵌入向量的点积得到一个(相似度)得分，然后传入Sigmoid函数，得到一个概率：
P
(
+
∣
w
t
,
c
)
=
σ
(
w
t
?
c
)
=
1
1
+
exp
?
(
?
w
t
?
c
)
(3)
P(+|w_t,c) = \sigma(w_t \cdot c) = \frac{1}{1 + \exp( - w_t \cdot c)} \tag{3}
P(+∣wt?,c)=σ(wt??c)=1+exp(?wt??c)1?(3)
同时我们要满足
(
11
)
(11)
(11)，即
c
c
c不是上下文单词的概率为：
P
(
?
∣
w
t
,
c
)
=
1
?
P
(
+
∣
w
t
,
c
)
=
σ
(
?
w
t
?
c
)
=
1
1
+
exp
?
(
w
?
c
)
(4)
\begin{aligned} P(-|w_t,c) &= 1 - P(+|w_t,c) \\ &= \sigma( - w_t \cdot c) = \frac{1}{1 + \exp( w \cdot c)} \end{aligned} \tag{4}
P(?∣wt?,c)?=1?P(+∣wt?,c)=σ(?wt??c)=1+exp(w?c)1??(4)
其中
1
?
1
1
+
exp
?
(
?
x
)
=
1
1
+
exp
?
(
x
)
1 - \frac{1}{1 + \exp(-x)} = \frac{1}{1 + \exp(x)}
1?1+exp(?x)1?=1+exp(x)1?很好证明，这里就不展开了。
这样我们得到了其中一个上下文单词的概率，但是窗口内包含很多个(
L
=
2
k
L=2k
L=2k)上下文单词。Skip-gram简化为所有上下文单词都是独立的假设，我们只需要让它们的概率相乘：
P
(
+
∣
w
t
,
c
1
:
L
)
=
∏
i
=
1
L
σ
(
w
t
?
c
i
)
(5)
P(+|w_t,c_{1:L}) = \prod_{i=1}^L \sigma( w_t \cdot c_i) \tag{5}
P(+∣wt?,c1:L?)=i=1∏L?σ(wt??ci?)(5)
我们使用取对数的基操，变成连加，防止数值溢出：
log
?
P
(
+
∣
w
t
,
c
1
:
L
)
=
∑
i
=
1
L
log
?
σ
(
w
t
?
c
i
)
(6)
\log P(+|w_t,c_{1:L}) = \sum_{i=1}^L \log \sigma( w_t \cdot c_i) \tag{6}
logP(+∣wt?,c1:L?)=i=1∑L?logσ(wt??ci?)(6)
其中
w
t
w_t
wt?和
c
i
c_i
ci?都表示词嵌入向量，计算方法在之前的Skip-gram模型中有介绍。
我们的模型定义好了，接下来看如何训练。
训练
如果我们想普通的Skip-gram模型一样，光有正例是不够的，那你的模型直接输出 1 1 1就好了。因此，我们需要负例，这也是负采样的由来。
我们需要让模型为正例尽可能输出 1 1 1，为负例尽可能输出 0 0 0。
我们考虑一个简单的例子：
... I love natural language processing ...
这里假设窗口大小
k
=
2
k=2
k=2，有一个目标词natural和
4
4
4个上下文单词，我们可以得到
4
4
4个正例：
| w w w | c p o s c_{pos} cpos? |
|---|---|
| natural | I |
| natural | love |
| natural | language |
| natural | processing |
这里 w w w表示目标词， c p o s c_{pos} cpos?表示真正的上下文单词， ( w , c p o s ) (w,c_{pos}) (w,cpos?)组成正例。上面说到，我们也需要负例。实际上SGNS使用了比正例数量更多的负例(有参数 K K K控制)。对于上面的每个正例，我们创建 K K K个负例，每个包含目标词和一个随机噪声单词。
噪声单词从词典中随机采样，但不能是上下文单词。这里的采样有一定的技巧。
使用加权unigram频率 P α ( w ) P_\alpha(w) Pα?(w)采样，其中 α \alpha α是一个权重。
那为什么需要加权呢？
我们对比下未加权的方法和加权的unigram。假设我们根据未加权频率
P
(
w
)
P(w)
P(w)进行采样，假设一个很罕见的单词aardvark，其概率
P
(
a
a
r
d
v
a
r
k
)
=
0.01
P(aardvark)=0.01
P(aardvark)=0.01。为了看到效果，夸张一点，假设另一个单词the出现的概率
P
(
t
h
e
)
=
0.99
P(the)=0.99
P(the)=0.99。
未加权说的是，我们只有
1
%
1\%
1%的概率抽取到单词aardvark。
再看下加权的情况，一般令
α
=
0.75
\alpha=0.75
α=0.75。那么有：
P
α
(
w
)
=
f
(
w
)
α
∑
w
′
f
(
w
′
)
α
(7)
P_\alpha(w) = \frac{\text{f}(w)^\alpha}{\sum_{w^\prime} \text{f}(w^\prime)^\alpha} \tag{7}
Pα?(w)=∑w′?f(w′)αf(w)α?(7)
我们计算：
P
α
(
t
h
e
)
=
0.9
9
0.75
0.9
9
0.75
+
0.
1
0.75
=
0.97
P
α
(
a
a
r
d
v
a
r
k
)
=
0.0
1
0.75
0.9
9
0.75
+
0.
1
0.75
=
0.03
\begin{aligned} P_\alpha(the) = \frac{0.99^{0.75}}{0.99^{0.75} + 0.1^{0.75}} = 0.97 \\ P_\alpha(aardvark) = \frac{0.01^{0.75}}{0.99^{0.75} + 0.1^{0.75}} = 0.03 \\ \end{aligned}
Pα?(the)=0.990.75+0.10.750.990.75?=0.97Pα?(aardvark)=0.990.75+0.10.750.010.75?=0.03?
这样，有更高的概率采样罕见单词。
假设我们令 K = 2 K=2 K=2，即对于每个正例，我们采样 2 2 2个负例，假设采样的负例为：
| w w w | c n e g c_{neg} cneg? |
|---|---|
| natural | where |
| natural | if |
| natural | jam |
| natural | ping |
| natural | coaxial |
| natural | oh |
| natural | pang |
| natural | bang |
由于我们有 4 4 4个正例，我们就采样了 8 8 8个负例。
那么目标就是训练这个分类器，使得
- 最大化正例中目标词和上下文词对 ( w , c p o s ) (w,c_{pos}) (w,cpos?)出现的概率
- 最小化负例中 ( w , c n e g ) (w,c_{neg}) (w,cneg?)词对出现的概率
假设我们考虑一个目标词/上下文词对
(
w
,
c
p
o
s
)
(w,c_{pos})
(w,cpos?)和
K
K
K个噪音单词
c
n
e
g
1
,
?
?
,
c
n
e
g
K
c_{neg_1},\cdots,c_{neg_K}
cneg1??,?,cnegK??，那么基于公式
(
15
)
(15)
(15)，我们需要最小化损失(所以加了个负号)，并同时考虑这两个目标：
L
C
E
=
?
log
?
[
P
(
+
∣
w
,
c
p
o
s
)
∏
i
=
1
K
P
(
?
∣
w
,
c
n
e
g
i
)
]
=
?
[
log
?
P
(
+
∣
w
,
c
p
o
s
)
+
∑
i
=
1
K
log
?
P
(
?
∣
w
,
c
n
e
g
i
)
]
=
?
[
log
?
P
(
+
∣
w
,
c
p
o
s
)
+
∑
i
=
1
K
log
?
(
1
?
P
(
+
∣
w
,
c
n
e
g
i
)
)
]
=
?
[
log
?
σ
(
c
p
o
s
?
w
)
+
∑
i
=
1
K
log
?
σ
(
?
c
n
e
g
i
?
w
)
]
(8)
\begin{aligned} L_{CE} &= - \log \left[ P(+|w,c_{pos}) \prod_{i=1}^K P(-|w,c_{neg_i}) \right] \\ &= -\left[ \log P(+|w,c_{pos}) + \sum_{i=1}^K \log P(-|w,c_{neg_i}) \right]\\ &= -\left[ \log P(+|w,c_{pos}) + \sum_{i=1}^K \log \left(1-P(+|w,c_{neg_i}) \right) \right] \\ &= -\left[ \log \sigma(c_{pos} \cdot w) + \sum_{i=1}^K \log \sigma( -c_{neg_i} \cdot w ) \right] \\ \end{aligned} \tag{8}
LCE??=?log[P(+∣w,cpos?)i=1∏K?P(?∣w,cnegi??)]=?[logP(+∣w,cpos?)+i=1∑K?logP(?∣w,cnegi??)]=?[logP(+∣w,cpos?)+i=1∑K?log(1?P(+∣w,cnegi??))]=?[logσ(cpos??w)+i=1∑K?logσ(?cnegi???w)]?(8)
这里分为两项，第一项我们希望分类器给正例很高的概率判断为
+
+
+；第二项希望给负例很高的概率判断为
?
-
?。
Sigmoid函数中计算了点积，我们想要最大化目标词与真正上下文单词的点积，同时最小化目标词与 K K K个负样本的点积。
到此为止就已经可以实现模型了，但是我们深入一步，推导一下对每种嵌入的梯度。
首先看对
c
p
o
s
c_{pos}
cpos?的梯度：
?
L
C
E
?
c
p
o
s
=
?
?
[
log
?
σ
(
c
p
o
s
?
w
)
+
∑
i
=
1
K
log
?
σ
(
?
c
n
e
g
i
?
w
)
]
?
c
p
o
s
=
?
σ
(
c
p
o
s
?
w
)
′
?
w
σ
(
c
p
o
s
?
w
)
=
?
σ
(
c
p
o
s
?
w
)
[
1
?
σ
(
c
p
o
s
?
w
)
]
?
w
σ
(
c
p
o
s
?
w
)
=
?
[
1
?
σ
(
c
p
o
s
?
w
)
]
?
w
=
[
σ
(
c
p
o
s
?
w
)
?
1
]
?
w
(9)
\begin{aligned} \frac{\partial L_{CE}}{\partial c_{pos}} &= - \frac{ \partial\left[ \log \sigma(c_{pos} \cdot w) + \sum_{i=1}^K \log \sigma( -c_{neg_i} \cdot w ) \right]}{\partial c_{pos}} \\ &= - \frac{\sigma(c_{pos} \cdot w)^\prime \cdot w}{\sigma(c_{pos} \cdot w)} \\ &= - \frac{\sigma(c_{pos} \cdot w) [1- \sigma(c_{pos} \cdot w)] \cdot w}{\sigma(c_{pos} \cdot w)} \\ &= -[1 - \sigma(c_{pos} \cdot w)] \cdot w \\ &= [\sigma(c_{pos} \cdot w) - 1]\cdot w \end{aligned} \tag{9}
?cpos??LCE???=??cpos??[logσ(cpos??w)+∑i=1K?logσ(?cnegi???w)]?=?σ(cpos??w)σ(cpos??w)′?w?=?σ(cpos??w)σ(cpos??w)[1?σ(cpos??w)]?w?=?[1?σ(cpos??w)]?w=[σ(cpos??w)?1]?w?(9)
然后是对第
i
i
i个负样本
c
n
e
g
i
c_{neg_i}
cnegi??的梯度：
?
L
C
E
?
c
n
e
g
i
=
?
?
[
log
?
σ
(
c
p
o
s
?
w
)
+
∑
i
=
1
K
log
?
σ
(
?
c
n
e
g
i
?
w
)
]
?
c
n
e
g
i
=
?
σ
(
?
c
n
e
g
i
?
w
)
′
?
(
?
w
)
σ
(
?
c
n
e
g
i
?
w
)
=
?
σ
(
?
c
n
e
g
i
?
w
)
[
1
?
σ
(
?
c
n
e
g
i
?
w
)
]
?
(
?
w
)
σ
(
?
c
n
e
g
i
?
w
)
=
σ
(
c
n
e
g
i
?
w
)
?
w
(10)
\begin{aligned} \frac{\partial L_{CE}}{\partial c_{neg_i}} &= - \frac{ \partial\left[ \log \sigma(c_{pos} \cdot w) + \sum_{i=1}^K \log \sigma( -c_{neg_i} \cdot w ) \right]}{\partial c_{neg_i}} \\ &= - \frac{\sigma( -c_{neg_i} \cdot w )^\prime \cdot (-w)}{\sigma( -c_{neg_i} \cdot w )} \\ &= - \frac{\sigma( -c_{neg_i} \cdot w )[1 - \sigma( -c_{neg_i} \cdot w )] \cdot (-w)}{\sigma( -c_{neg_i} \cdot w )} \\ &= \sigma( c_{neg_i} \cdot w ) \cdot w \end{aligned} \tag{10}
?cnegi???LCE???=??cnegi???[logσ(cpos??w)+∑i=1K?logσ(?cnegi???w)]?=?σ(?cnegi???w)σ(?cnegi???w)′?(?w)?=?σ(?cnegi???w)σ(?cnegi???w)[1?σ(?cnegi???w)]?(?w)?=σ(cnegi???w)?w?(10)
最后看一下对中心词
w
w
w的梯度：
?
L
C
E
?
w
=
?
?
[
log
?
σ
(
c
p
o
s
?
w
)
+
∑
i
=
1
K
log
?
σ
(
?
c
n
e
g
i
?
w
)
]
?
w
=
?
σ
(
c
p
o
s
?
w
)
′
?
c
p
o
s
σ
(
c
p
o
s
?
w
)
?
∑
i
=
1
K
[
σ
(
?
c
n
e
g
i
?
w
)
′
?
(
?
c
n
e
g
i
)
σ
(
?
c
n
e
g
i
?
w
)
]
=
?
σ
(
c
p
o
s
?
w
)
[
1
?
σ
(
c
p
o
s
?
w
)
]
?
c
p
o
s
σ
(
c
p
o
s
?
w
)
?
∑
i
=
1
K
[
σ
(
?
c
n
e
g
i
?
w
)
[
1
?
σ
(
?
c
n
e
g
i
?
w
)
]
?
(
?
c
n
e
g
i
)
σ
(
?
c
n
e
g
i
?
w
)
]
=
?
[
1
?
σ
(
c
p
o
s
?
w
)
]
?
c
p
o
s
+
∑
i
=
1
K
σ
(
c
n
e
g
i
?
w
)
?
c
n
e
g
i
=
[
σ
(
c
p
o
s
?
w
)
?
1
]
?
c
p
o
s
+
∑
i
=
1
K
σ
(
c
n
e
g
i
?
w
)
?
c
n
e
g
i
(11)
\begin{aligned} \frac{\partial L_{CE}}{\partial w} &= - \frac{ \partial\left[ \log \sigma(c_{pos} \cdot w) + \sum_{i=1}^K \log \sigma( -c_{neg_i} \cdot w ) \right]}{\partial w} \\ &= - \frac{ \sigma(c_{pos} \cdot w)^\prime \cdot c_{pos}}{ \sigma(c_{pos} \cdot w)} - \sum_{i=1}^K \left[ \frac{\sigma( -c_{neg_i} \cdot w ) ^\prime \cdot (- c_{neg_i})}{\sigma( -c_{neg_i} \cdot w ) } \right] \\ &= - \frac{ \sigma(c_{pos} \cdot w)[1 - \sigma(c_{pos} \cdot w)] \cdot c_{pos}}{ \sigma(c_{pos} \cdot w)} - \sum_{i=1}^K \left[ \frac{\sigma( -c_{neg_i} \cdot w )[1- \sigma( -c_{neg_i} \cdot w )]\cdot (- c_{neg_i})}{\sigma( -c_{neg_i} \cdot w ) } \right] \\ &= - [1 - \sigma(c_{pos} \cdot w)] \cdot c_{pos} + \sum_{i=1}^K \sigma( c_{neg_i} \cdot w ) \cdot c_{neg_i} \\ &= [ \sigma(c_{pos} \cdot w) -1]\cdot c_{pos} + \sum_{i=1}^K \sigma( c_{neg_i} \cdot w ) \cdot c_{neg_i} \end{aligned} \tag{11}
?w?LCE???=??w?[logσ(cpos??w)+∑i=1K?logσ(?cnegi???w)]?=?σ(cpos??w)σ(cpos??w)′?cpos???i=1∑K?[σ(?cnegi???w)σ(?cnegi???w)′?(?cnegi??)?]=?σ(cpos??w)σ(cpos??w)[1?σ(cpos??w)]?cpos???i=1∑K?[σ(?cnegi???w)σ(?cnegi???w)[1?σ(?cnegi???w)]?(?cnegi??)?]=?[1?σ(cpos??w)]?cpos?+i=1∑K?σ(cnegi???w)?cnegi??=[σ(cpos??w)?1]?cpos?+i=1∑K?σ(cnegi???w)?cnegi???(11)
因为
w
w
w参与了这两项，所以它的式子也由两项组成。
代码实现
首先构建SGNS数据集，对于每个训练(正)样本，需要根据某个负采样概率分布生成相应的负样本，同时需要保证负样本不包含当前上下文中的词。
一种实现方式是，在构建训练数据的过程中就完成负样本的生成，这样在训练时直接读取负样本即可。这么做的优点是训练过程无需再进行采样，因此效率较高；缺点是每次迭代使用的是同样的负样本，缺乏多样性。
这里采用在训练过程中实时进行负采样的实现方式，通过以下类的collate_fn函数完成负采样。
class SGNSDataset(Dataset):
def __init__(self, corpus, vocab, window_size=2, n_negatives=5, ns_dist=None):
self.data = []
self.bos = vocab[BOS_TOKEN]
self.eos = vocab[EOS_TOKEN]
self.pad = vocab[PAD_TOKEN]
for sentence in tqdm(corpus, desc='Dataset Construction'):
sentence = [self.bos] + sentence + [self.eos]
for i in range(1, len(sentence) - 1):
# 模型输入：(w, context)
# 输出：0/1，表示context是否为负样本
w = sentence[i]
left_context_index = max(0, i - window_size)
right_context_index = min(len(sentence), i + window_size)
context = sentence[left_context_index:i] + sentence[i + 1:right_context_index + 1]
context += [self.pad] * (2 * window_size - len(context))
self.data.append((w, context))
# 负样本数量
self.n_negatives = n_negatives
# 负采样分布：若参数ns_dist为None，则使用uniform分布
self.ns_dist = ns_dist if ns_dist is not None else Tensor.ones(len(vocab))
self.data = np.asarray(self.data)
def __len__(self):
return len(self.data)
def __getitem__(self, i):
return self.data[i]
def collate_fn(self, examples):
words = Tensor([ex[0] for ex in examples])
contexts = Tensor([ex[1] for ex in examples])
batch_size, window_size = contexts.shape
neg_contexts = []
# 对batch内的样本分别进行负采样
for i in range(batch_size):
# 保证负样本不包含当前样本中的context
ns_dist = self.ns_dist.index_fill_(0, contexts[i], .0)
neg_contexts.append(Tensor.multinomial(ns_dist, self.n_negatives * window_size, replace=True))
neg_contexts = F.stack(neg_contexts, axis=0)
return words, contexts, neg_contexts
在模型类中需要维护两个词向量w_embeddings和c_embeddings，分别用于词与上下文的向量表示。 同时因为word2vec模型比较特殊，我们不关心模型的输出，而是它学到的权重。为了简单，我们这里在forward中直接输出损失，使用公式
(
8
)
(8)
(8)来计算总损失：
class SGNSModel(nn.Module):
def __init__(self, vocab_size, embedding_dim):
# 目标词嵌入
self.w_embeddings = nn.Embedding(vocab_size, embedding_dim)
# 上下文嵌入
self.c_embeddings = nn.Embedding(vocab_size, embedding_dim)
def forward(self, target_words, pos_contexts, neg_contexts) -> Tensor:
'''
word2vec模型比较特殊，我们不关心模型的输出，而是它学到的权重
为了简单，我们这里直接输出损失
'''
batch_size = target_words.shape[0]
n_negatives = neg_contexts.shape[-1]
word_embeds = self.w_embeddings(target_words) # (batch_size, embedding_dim)
context_embeds = self.c_embeddings(pos_contexts) # (batch_size, window_size * 2, embedding_dim)
neg_context_embeds = self.c_embeddings(neg_contexts) # (batch_size, window_size * n_negatives, embedding_dim)
word_embeds = word_embeds.unsqueeze(2)
# 正样本的对数似然
context_loss = F.logsigmoid((context_embeds @ word_embeds).squeeze(2))
context_loss = context_loss.mean(axis=1)
# 负样本的对数似然
neg_context_loss = F.logsigmoid((neg_context_embeds @ word_embeds).squeeze(axis=2).neg())
neg_context_loss = neg_context_loss.reshape((batch_size, -1, n_negatives)).sum(axis=2)
neg_context_loss = neg_context_loss.mean(axis=1)
# 总损失： 负对数似然
loss = -(context_loss + neg_context_loss).mean()
return loss
但我们还需要编写从训练语料库中统计Unigram出现次数，并计算概率分布。以此概率为基础进行负采样：
def get_unigram_distribution(corpus, vocab_size):
# 从给定语料中统计unigram概率分布
token_counts = Tensor([.0] * vocab_size)
total_count = .0
for sentence in corpus:
total_count += len(sentence)
for token in sentence:
token_counts[token] += 1
unigram_dist = token_counts / total_count
return unigram_dist
下面是具体的训练过程：
embedding_dim = 64
window_size = 2
batch_size = 10240
num_epoch = 10
min_freq = 3 # 保留单词最少出现的次数
n_negatives = 10 # 负采样数
# 读取数据
corpus, vocab = load_corpus('../data/xiyouji.txt', min_freq)
# 计算unigram概率分布
unigram_dist = get_unigram_distribution(corpus, len(vocab))
# 根据unigram分布计算负采样分数： p(w)**0.75
negative_sampling_dist = unigram_dist ** 0.75
# 构建数据集
dataset = SGNSDataset(corpus, vocab, window_size=window_size, ns_dist=negative_sampling_dist)
# 构建数据加载器
data_loader = DataLoader(
dataset,
batch_size=batch_size,
collate_fn=dataset.collate_fn,
shuffle=True
)
device = cuda.get_device("cuda:0" if cuda.is_available() else "cpu")
print(f'current device:{device}')
# 构建模型
model = SGNSModel(len(vocab), embedding_dim)
model.to(device)
optimizer = SGD(model.parameters())
with debug_mode():
for epoch in range(num_epoch):
total_loss = 0
for batch in tqdm(data_loader, desc=f'Training Epoch {epoch}'):
words, contexts, neg_contexts = [x.to(device) for x in batch]
optimizer.zero_grad()
loss = model(words, contexts, neg_contexts)
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f'Loss: {total_loss:.2f}')
save_pretrained(vocab, model.embeddings.weight, 'sgns.vec')
完整代码
https://github.com/nlp-greyfoss/metagrad
References
- 从零实现Word2Vec
- 自然语言处理：基于预训练模型的方法
- Speech and Language Processing
