word2vec-gluon.md 16.6 KB
Newer Older
A
Aston Zhang 已提交
1
# Word2vec的实现
M
muli 已提交
2

A
Aston Zhang 已提交
3
本节是对前两节内容的实践。我们以[“词嵌入(word2vec)”](word2vec.md)一节中的跳字模型和[“近似训练”](approx-training.md)一节中的负采样为例,介绍在语料库上训练词嵌入模型的实现。我们还会介绍一些实现中的技巧,例如二次采样(subsampling)和掩码(mask)变量。
M
muli 已提交
4

A
Aston Zhang 已提交
5
首先让我们导入实验所需的包或模块。
M
muli 已提交
6

A
Aston Zhang 已提交
7
```{.python .input  n=1}
A
Aston Zhang 已提交
8 9 10
import sys
sys.path.insert(0, '..')

M
muli 已提交
11 12 13 14 15 16 17 18 19 20 21 22 23
import collections
import gluonbook as gb
import math
from mxnet import autograd, gluon, nd
from mxnet.gluon import data as gdata, loss as gloss, nn
import random
import sys
import time
import zipfile
```

## 处理数据集

A
Aston Zhang 已提交
24
Penn Tree Bank(PTB)是一个常用的小型语料库 [1]。它采样自华尔街日报的文章,包括训练集、验证集和测试集。我们将在PTB的训练集上训练词嵌入模型。该数据集的每一行作为一个句子。句子中的每个词由空格隔开。
M
muli 已提交
25

A
Aston Zhang 已提交
26
```{.python .input  n=2}
M
muli 已提交
27
with zipfile.ZipFile('../data/ptb.zip', 'r') as zin:
A
Aston Zhang 已提交
28 29 30 31 32 33 34
    zin.extractall('../data/')

with open('../data/ptb/ptb.train.txt', 'r') as f:
    lines = f.readlines()
    # st 是 sentence 在循环中的缩写。
    raw_dataset = [st.split() for st in lines]

M
muli 已提交
35 36 37
'# sentences: %d' % len(raw_dataset)
```

A
Aston Zhang 已提交
38
对于数据集的前三个句子,打印每个句子的词数和前五个词。这个数据集中句尾符为“<eos>”,生僻词全用“<unk>”表示,数字则被替换成了“N”。
M
muli 已提交
39

A
Aston Zhang 已提交
40
```{.python .input  n=3}
M
muli 已提交
41 42 43 44 45 46 47 48
for st in raw_dataset[:3]:
    print('# tokens:', len(st), st[:5])
```

### 建立词语索引

为了计算简单,我们只保留在数据集中至少出现5次的词。

A
Aston Zhang 已提交
49
```{.python .input  n=4}
M
muli 已提交
50 51 52 53 54 55 56
# tk 是 token 的在循环中的缩写。
counter = collections.Counter([tk for st in raw_dataset for tk in st])
counter = dict(filter(lambda x: x[1] >= 5, counter.items()))
```

然后将词映射到整数索引。

A
Aston Zhang 已提交
57
```{.python .input  n=5}
M
muli 已提交
58 59 60 61 62
idx_to_token = [tk for tk, _ in counter.items()]
token_to_idx = {tk: idx for idx, tk in enumerate(idx_to_token)}
dataset = [[token_to_idx[tk] for tk in st if tk in token_to_idx] 
           for st in raw_dataset]
num_tokens = sum([len(st) for st in dataset])
A
Aston Zhang 已提交
63
'# tokens: %d' % num_tokens
M
muli 已提交
64 65 66 67
```

### 二次采样

A
Aston Zhang 已提交
68
文本数据中一般会出现一些高频词,例如英文中的“the”、“a”和“in”。通常来说,在一个背景窗口中,一个词(如“chip”)和较低频词(如“microprocessor”)同时出现比和较高频词(如“the”)同时出现对训练词嵌入模型更有益。因此,训练词嵌入模型时可以对词进行二次采样 [2]。具体来说,数据集中每个被索引词$w_i$将有一定概率被丢弃,该丢弃概率为
M
muli 已提交
69

A
Aston Zhang 已提交
70
$$ \mathbb{P}(w_i) = \max\left(1 - \sqrt{\frac{t}{f(w_i)}}, 0\right),$$ 
M
muli 已提交
71

A
Aston Zhang 已提交
72
其中 $f(w_i)$ 是数据集中词$w_i$的个数与总词数之比,常数$t$是一个超参数(实验中设为$10^{-4}$)。可见,只有当$f(w_i) > t$时,我们才有可能在二次采样中丢弃词$w_i$,并且越高频的词被丢弃的概率越大。
M
muli 已提交
73

A
Aston Zhang 已提交
74
```{.python .input  n=6}
A
Aston Zhang 已提交
75 76 77 78
def discard(idx):
    return random.uniform(0, 1) < 1 - math.sqrt(
        1e-4 / counter[idx_to_token[idx]] * num_tokens)

M
muli 已提交
79
subsampled_dataset = [[tk for tk in st if not discard(tk)] for st in dataset]
A
Aston Zhang 已提交
80
'# tokens: %d' % sum([len(st) for st in subsampled_dataset])
M
muli 已提交
81 82
```

A
Aston Zhang 已提交
83
可以看到,二次采样后我们去掉了一半左右的词。下面比较一个词在二次采样前后出现在数据集中的次数。可见高频词“the”的采样率不足1/20。
M
muli 已提交
84

A
Aston Zhang 已提交
85
```{.python .input  n=7}
A
Aston Zhang 已提交
86 87
def compare_counts(token):
    return '# %s: before=%d, after=%d' % (token, sum(
A
Aston Zhang 已提交
88 89 90
        [st.count(token_to_idx[token]) for st in dataset]), sum(
        [st.count(token_to_idx[token]) for st in subsampled_dataset]))

A
Aston Zhang 已提交
91
compare_counts('the')
M
muli 已提交
92 93
```

A
Aston Zhang 已提交
94
但低频词“join”则完整地保留了下来。
M
muli 已提交
95

A
Aston Zhang 已提交
96
```{.python .input  n=8}
A
Aston Zhang 已提交
97
compare_counts('join')
M
muli 已提交
98 99 100 101
```

### 提取中心词和背景词

A
Aston Zhang 已提交
102
我们将与中心词距离不超过背景窗口大小的词作为它的背景词。下面定义函数提取出所有中心词和它们的背景词。它每次在整数1和`max_window_size`(最大背景窗口)之间均匀随机采样一个整数作为背景窗口大小。
M
muli 已提交
103

A
Aston Zhang 已提交
104
```{.python .input  n=9}
M
muli 已提交
105 106 107 108 109 110
def get_centers_and_contexts(dataset, max_window_size):
    centers, contexts = [], []
    for st in dataset:
        if len(st) < 2:  # 每个句子至少要有 2 个词才可能组成一对“中心词-背景词”。
            continue
        centers += st
A
Aston Zhang 已提交
111
        for center_i in range(len(st)):
M
muli 已提交
112
            window_size = random.randint(1, max_window_size)
A
Aston Zhang 已提交
113 114 115
            indices = list(range(max(0, center_i - window_size),
                                 min(len(st), center_i + 1 + window_size)))
            indices.remove(center_i)  # 将中心词排除在背景词之外。
M
muli 已提交
116 117 118 119
            contexts.append([st[idx] for idx in indices])
    return centers, contexts
```

A
Aston Zhang 已提交
120
下面我们创建一个人工数据集,其中含有词数分别为7和3的两个句子。设最大背景窗口为2,打印所有中心词和它们的背景词。
M
muli 已提交
121

A
Aston Zhang 已提交
122
```{.python .input  n=10}
A
Aston Zhang 已提交
123 124
tiny_dataset = [list(range(7)), list(range(7, 10))]
print('dataset', tiny_dataset)
M
muli 已提交
125 126 127 128
for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)):
    print('center', center, 'has contexts', context)
```

A
Aston Zhang 已提交
129
实验中,我们设最大背景窗口大小为5。下面提取数据集中所有的中心词及其背景词。
M
muli 已提交
130

A
Aston Zhang 已提交
131
```{.python .input  n=11}
M
muli 已提交
132 133 134 135 136
all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5)
```

## 负采样

A
Aston Zhang 已提交
137
我们使用负采样来进行近似训练。对于一对中心词和背景词,我们随机采样$K$个噪音词(实验中设$K=5$)。根据word2vec论文的建议,噪音词采样概率$\mathbb{P}(w)$设为$w$词频与总词频之比的0.75次方 [2]。
M
muli 已提交
138

A
Aston Zhang 已提交
139
```{.python .input  n=12}
A
Aston Zhang 已提交
140 141 142 143
def get_negatives(all_contexts, sampling_weights, K):
    all_negatives, neg_candidates, i = [], [], 0
    population = list(range(len(sampling_weights)))
    for contexts in all_contexts:
M
muli 已提交
144 145
        negatives = []
        while len(negatives) < len(contexts) * K:
A
Aston Zhang 已提交
146 147 148 149 150 151 152 153
            if i == len(neg_candidates):
                # 根据每个词的权重(sampling_weights)随机生成 k 个词的索引作为噪音
                # 词。为了高效计算,可以将 k 设的稍大一点。
                i, neg_candidates = 0, random.choices(
                    population, sampling_weights, k=int(1e5))
            neg, i = neg_candidates[i], i + 1
            # 噪音词不能是背景词。
            if neg not in set(contexts):
M
muli 已提交
154 155 156
                negatives.append(neg)
        all_negatives.append(negatives)
    return all_negatives
A
Aston Zhang 已提交
157 158 159

sampling_weights = [counter[w]**0.75 for w in idx_to_token]
all_negatives = get_negatives(all_contexts, sampling_weights, 5)
M
muli 已提交
160 161
```

A
Aston Zhang 已提交
162
## 读取数据
M
muli 已提交
163

A
Aston Zhang 已提交
164
我们从数据集中提取所有中心词`all_centers`,以及每个中心词对应的背景词`all_contexts`和噪音词`all_negatives`。我们将通过随机小批量来读取它们。
M
muli 已提交
165

A
Aston Zhang 已提交
166
在一个小批量数据中,第$i$个样本包括一个中心词以及它所对应的$n_i$个背景词和$m_i$个噪音词。由于每个样本的背景窗口大小可能不一样,其中背景词与噪音词个数之和$n_i+m_i$也会不同。在构造小批量时,我们将每个样本的背景词和噪音词连结在一起,并添加填充项0直至连结后的长度相同,即长度均为$\max_i n_i+m_i$(`max_len`)。为了避免填充项对损失函数计算的影响,我们构造了掩码变量`masks`,其每一个元素分别与连结后的背景词和噪音词`contexts_negatives`中的元素一一对应。当变量`contexts_negatives`中的某个元素为填充项时,相同位置的掩码变量`masks`中的元素取0,否则取1。为了区分正例和负例,我们还需要将`contexts_negatives`变量中的背景词和噪音词区分开来。依据掩码变量的构造思路,我们只需创建与`contexts_negatives`变量形状相同的标签变量`labels`,并将与背景词(正例)对应的元素设1,其余清0。
M
muli 已提交
167

A
Aston Zhang 已提交
168
下面我们将实现这个小批量读取函数`batchify`。它的小批量输入`data`是一个长度为批量大小的列表,其中每个元素分别包含中心词`center`、背景词`context`和噪音词`negative`。该函数返回的小批量数据符合我们所需要的格式,例如包含了掩码变量。
M
muli 已提交
169

A
Aston Zhang 已提交
170
```{.python .input  n=13}
M
muli 已提交
171
def batchify(data):
A
Aston Zhang 已提交
172
    max_len = max(len(c) + len(n) for _, c, n in data)
M
muli 已提交
173 174
    centers, contexts_negatives, masks, labels = [], [], [], []
    for center, context, negative in data:
A
Aston Zhang 已提交
175
        cur_len = len(context) + len(negative)
M
muli 已提交
176
        centers += [center]
A
Aston Zhang 已提交
177 178 179
        contexts_negatives += [context + negative + [0] * (max_len - cur_len)]
        masks += [[1] * cur_len + [0] * (max_len - cur_len)]
        labels += [[1] * len(context) + [0] * (max_len - len(context))]
M
muli 已提交
180 181 182 183
    return (nd.array(centers).reshape((-1, 1)), nd.array(contexts_negatives),
            nd.array(masks), nd.array(labels))
```

A
Aston Zhang 已提交
184
我们用刚刚定义的`batchify`函数指定`DataLoader`实例中小批量的读取方式。然后打印读取的第一个批量中各个变量的形状。
M
muli 已提交
185

A
Aston Zhang 已提交
186
```{.python .input  n=14}
M
muli 已提交
187 188 189 190 191 192
batch_size = 512
num_workers = 0 if sys.platform.startswith('win32') else 4
dataset = gdata.ArrayDataset(all_centers, all_contexts, all_negatives)
data_iter = gdata.DataLoader(dataset, batch_size, shuffle=True,
                             batchify_fn=batchify, num_workers=num_workers)
for batch in data_iter:
A
Aston Zhang 已提交
193 194 195
    for name, data in zip(['centers', 'contexts_negatives', 'masks',
                           'labels'], batch):
        print(name, 'shape:', data.shape)
M
muli 已提交
196 197 198 199 200 201 202 203 204 205 206
    break
```

## 跳字模型

给定中心词$w_c$,和其对应的$n$个背景词$w_{o_1}, \ldots, w_{o_n}$,以及采样得到的$m$个噪音词$w_{o_{n+1}}, \ldots, w_{o_{n+m}}$。在跳字模型我们首先要得到它们对应的词嵌入$\boldsymbol{v}_c$, $\boldsymbol{u}_{o_1}, \ldots, \boldsymbol{u}_{o_n}$, $\boldsymbol{u}_{o_{n+1}}, \ldots, \boldsymbol{u}_{o_{n+m}}$。

### 嵌入层

获取词嵌入的层被称为嵌入层。嵌入层的权重是一个(词典大小,嵌入向量长度$d$)的矩阵。输入一个词的索引$i$,它返回权重的第$i$行作为它的词嵌入向量。在Gluon中可以通过`nn.Embedding`类来得到嵌入层。

A
Aston Zhang 已提交
207
```{.python .input  n=15}
M
muli 已提交
208 209 210 211 212 213 214
embed = nn.Embedding(input_dim=20, output_dim=2)
embed.initialize()
embed.weight
```

嵌入层输入为词索引。假设输入形状为(批量大小,$l$),那么输出的形状为(批量大小,$l$,$d$),最后一个维度用来放置嵌入向量。

A
Aston Zhang 已提交
215
```{.python .input  n=16}
M
muli 已提交
216 217 218 219 220 221
x = nd.array([[1,2,3]])
embed(x)
```

### 小批量乘法

究其根本's avatar
究其根本 已提交
222
我们可以将背景词向量和噪音词向量合并起来,然后使用一次矩阵乘法来计算中心词向量和它们的内积$\boldsymbol{v}_{c}^\top\left[\boldsymbol{u}_{o_1},\ldots,\boldsymbol{u}_{o_{n+m}}\right]$。我们需要对小批量里的每个中心词逐一做此运算,虽然可以用for循环来实现,但使用`batch_dot`(用$\text{bd}$表示)通常可以得到更好的性能。假设$\boldsymbol{X}=\left[\boldsymbol{X}_1,\ldots,\boldsymbol{X}_n\right]$且$\boldsymbol{Y}=\left[\boldsymbol{Y}_1,\ldots,\boldsymbol{Y}_n\right]$,如果$\boldsymbol{Z}=\text{bd}(\boldsymbol{X},\boldsymbol{Y})$,那么$\boldsymbol{Z}=\left[\boldsymbol{X}_1\boldsymbol{Y}_1,\ldots,\boldsymbol{X}_n\boldsymbol{Y}_n\right]$。
M
muli 已提交
223

A
Aston Zhang 已提交
224
```{.python .input  n=17}
M
muli 已提交
225 226 227 228 229 230 231 232 233
X = nd.ones((2, 3, 4))
Y = nd.ones((2, 4, 6))
nd.batch_dot(X, Y).shape
```

### 跳字模型前向计算

在前向计算中,跳字模型的输入是形状为(批量大小,1)的中心词索引,形状为(批量大小,$n+m$)的背景词和噪音词索引,以及对应的嵌入层。输出形状为(批量大小,1,$n+m$),每个元素是一对嵌入向量的内积。

A
Aston Zhang 已提交
234
```{.python .input  n=18}
M
muli 已提交
235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
def skip_gram(center, contexts_and_negatives, embed_v, embed_u):
    v = embed_v(center)
    u = embed_u(contexts_and_negatives)
    pred = nd.batch_dot(v, u.swapaxes(1, 2))
    return pred
```

## 模型训练

### 二元交叉熵损失函数

我们知道负采样的损失函数为:

$$\sum_{k=1}^n\log\,\sigma(\boldsymbol{u}_{o_k}^\top\boldsymbol{v}_c) + \sum_{k=n+1}^{n+m}\log\,\sigma(-\boldsymbol{u}_{o_k}^\top\boldsymbol{v}_c),$$

究其根本's avatar
究其根本 已提交
250
这里$\sigma(x)=1/(1+\exp(-x))$是sigmoid函数。这个等价于一个使用交叉熵损失函数的两类softmax回归,又称logistic回归。如果构造标签$y_1,\ldots, y_{n+m}$,使得 $y_1=\ldots=y_n=1$,$y_{n+1}=\ldots=y_{n+m}=0$,那么我们可以重写上面损失函数为标准的logistic回归目标函数:
M
muli 已提交
251 252 253 254 255

$$\sum_{k=1}^{n+m}y_k\log\,\sigma(\boldsymbol{u}_{o_k}^\top\boldsymbol{v}_c) + (1-y_k)\log\,\sigma(-\boldsymbol{u}_{o_k}^\top\boldsymbol{v}_c)$$

从而我们可以直接使用Gluon提供的`SigmoidBinaryCrossEntropyLoss`

A
Aston Zhang 已提交
256
```{.python .input  n=19}
M
muli 已提交
257 258 259 260 261 262 263
loss = gloss.SigmoidBinaryCrossEntropyLoss()
```

### 初始化模型参数

我们构造中心词和背景词对应的嵌入层,并将超参数嵌入向量长度设置成100。

A
Aston Zhang 已提交
264
```{.python .input  n=20}
M
muli 已提交
265 266 267 268 269 270 271 272 273 274
embed_size = 100
net = nn.Sequential()
net.add(nn.Embedding(input_dim=len(idx_to_token), output_dim=embed_size),
        nn.Embedding(input_dim=len(idx_to_token), output_dim=embed_size))
```

### 训练

下面定义训练函数。注意由于填充的关系,在计算损失处跟之前的训练函数略微不同。

A
Aston Zhang 已提交
275
```{.python .input  n=21}
M
muli 已提交
276 277 278
def train(net, lr, num_epochs):
    ctx = gb.try_gpu()
    net.initialize(ctx=ctx, force_reinit=True)
A
Aston Zhang 已提交
279 280
    trainer = gluon.Trainer(net.collect_params(), 'adam',
                            {'learning_rate': lr})
M
muli 已提交
281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301
    for epoch in range(num_epochs):
        start_time, train_l = time.time(), 0
        for batch in data_iter:
            center, context_negative, mask, label = [
                data.as_in_context(ctx) for data in batch]
            with autograd.record():
                pred = skip_gram(center, context_negative, net[0], net[1])
                # 将 mask 通过 loss 的权重选项来过滤掉填充。
                l = loss(pred.reshape(label.shape), label, mask)
                # l[i] 是第 i 个样本背景词和噪音词损失的均值。但由于填充关系,这个均值
                # 不正确。这里我们修改 l 为对小批量中所有背景词和噪音词的平均损失。
                l = l.sum() * mask.shape[1] / mask.sum()
            l.backward()
            trainer.step(1)  # l 已经做过了平均,这里不需要传入 batch_size。
            train_l += l.asscalar()
        print('epoch %d, train loss %.2f, time %.2fs' % (
            epoch + 1, train_l / len(data_iter), time.time() - start_time))
```

现在我们可以训练使用负采样的跳字模型了。可以看到,使用训练得到的词嵌入模型时,与词“chip”语义最接近的词大多与芯片有关。

A
Aston Zhang 已提交
302
```{.python .input  n=22}
M
muli 已提交
303 304 305 306 307 308 309
train(net, 0.005, 8)
```

## 模型预测

当我们训练好词嵌入模型后,我们可以根据两个词向量的余弦相似度表示词与词之间在语义上的相似度。

A
Aston Zhang 已提交
310
```{.python .input  n=23}
M
muli 已提交
311
def get_similar_tokens(query_token, k, embed):
M
muli 已提交
312
    W = embed.weight.data()
A
Aston Zhang 已提交
313 314
    x = W[token_to_idx[query_token]]
    cos = nd.dot(W, x) / nd.sum(W * W, axis=1).sqrt() / nd.sum(x * x).sqrt()
M
muli 已提交
315 316 317
    topk = nd.topk(cos, k=k+1, ret_typ='indices').asnumpy().astype('int32')
    for i in topk[1:]:  # 除去输入词。
        print('similarity=%.3f: %s' % (cos[i].asscalar(), (idx_to_token[i])))
A
Aston Zhang 已提交
318
get_similar_tokens('chip', 3, net[0])
M
muli 已提交
319 320 321 322 323 324 325 326 327 328 329 330 331
```

## 小结

* 我们使用Gluon通过负采样训练了跳字模型。
* 二次采样试图尽可能减轻高频词对训练词嵌入模型的影响。
* 我们可以将长度不同的样本填充至长度相同的小批量,并通过掩码变量区分非填充和填充,然后只令非填充参与损失函数的计算。


## 练习

* 调一调模型和训练超参数。
* 试着找出其他词的近义词。
究其根本's avatar
究其根本 已提交
332
* 当数据集较大时,我们通常在迭代模型参数时才对当前小批量里的中心词采样背景词和音词。也就是说,同一个中心词在不同的迭代周期可能会有不同的背景词或噪音词。这样训练有哪些好处?尝试实现上述训练方法。
M
muli 已提交
333 334 335 336 337
* 事实上`SigmoidBinaryCrossEntropyLoss`的实现不是本节给出的目标函数。找出它的实现,并分析为什么要这么做。


## 扫码直达[讨论区](https://discuss.gluon.ai/t/topic/7761)

A
add qr  
Aston Zhang 已提交
338
![](../img/qr_word2vec-gluon.svg)
M
muli 已提交
339 340 341 342 343 344 345


## 参考文献

[1] Penn Tree Bank. https://catalog.ldc.upenn.edu/ldc99t42

[2] Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., & Dean, J. (2013). Distributed representations of words and phrases and their compositionality. In Advances in neural information processing systems (pp. 3111-3119).