10.3 word2vec實現 · 動手學深度學習--pytorch

# 10.3 word2vec的實現本節是對前兩節內容的實踐。我們以10.1節（詞嵌入word2vec）中的跳字模型和10.2節（近似訓練）中的負采樣為例，介紹在語料庫上訓練詞嵌入模型的實現。我們還會介紹一些實現中的技巧，如二次采樣（subsampling）。首先導入實驗所需的包或模塊。 ``` python import collections import math import random import sys import time import os import numpy as np import torch from torch import nn import torch.utils.data as Data sys.path.append("..") import d2lzh_pytorch as d2l print(torch.__version__) ``` ## 10.3.1 處理數據集 PTB（Penn Tree Bank）是一個常用的小型語料庫 [1]。它采樣自《華爾街日報》的文章，包括訓練集、驗證集和測試集。我們將在PTB訓練集上訓練詞嵌入模型。該數據集的每一行作為一個句子。句子中的每個詞由空格隔開。確保`ptb.train.txt`已經放在了文件夾`../../data/ptb`下。 ``` python assert 'ptb.train.txt' in os.listdir("../../data/ptb") with open('../../data/ptb/ptb.train.txt', 'r') as f: lines = f.readlines() # st是sentence的縮寫 raw_dataset = [st.split() for st in lines] '# sentences: %d' % len(raw_dataset) # 輸出 '# sentences: 42068' ``` 對于數據集的前3個句子，打印每個句子的詞數和前5個詞。這個數據集中句尾符為"<eos>"，生僻詞全用"<unk>"表示，數字則被替換成了"N"。 ``` python for st in raw_dataset[:3]: print('# tokens:', len(st), st[:5]) ``` 輸出： ``` # tokens: 24 ['aer', 'banknote', 'berlitz', 'calloway', 'centrust'] # tokens: 15 ['pierre', '<unk>', 'N', 'years', 'old'] # tokens: 11 ['mr.', '<unk>', 'is', 'chairman', 'of'] ``` ### 10.3.1.1 建立詞語索引為了計算簡單，我們只保留在數據集中至少出現5次的詞。 ``` python # 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())) ``` 然后將詞映射到整數索引。 ``` python 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]) '# tokens: %d' % num_tokens # 輸出 '# tokens: 887100' ``` ### 10.3.1.2 二次采樣文本數據中一般會出現一些高頻詞，如英文中的“the”“a”和“in”。通常來說，在一個背景窗口中，一個詞（如“chip”）和較低頻詞（如“microprocessor”）同時出現比和較高頻詞（如“the”）同時出現對訓練詞嵌入模型更有益。因此，訓練詞嵌入模型時可以對詞進行二次采樣 [2]。具體來說，數據集中每個被索引詞$w_i$將有一定概率被丟棄，該丟棄概率為 ```[tex] P(w_i) = \max\left(1 - \sqrt{\frac{t}{f(w_i)}}, 0\right), ``` 其中 `$ f(w_i) $`是數據集中詞`$ w_i $`的個數與總詞數之比，常數`$ t $是`一個超參數（實驗中設為`$ 10^{-4} $`）。可見，只有當`$ f(w_i) > t $`時，我們才有可能在二次采樣中丟棄詞`$ w_i $`，并且越高頻的詞被丟棄的概率越大。 ``` python def discard(idx): return random.uniform(0, 1) < 1 - math.sqrt( 1e-4 / counter[idx_to_token[idx]] * num_tokens) subsampled_dataset = [[tk for tk in st if not discard(tk)] for st in dataset] '# tokens: %d' % sum([len(st) for st in subsampled_dataset]) # '# tokens: 375875' ``` 可以看到，二次采樣后我們去掉了一半左右的詞。下面比較一個詞在二次采樣前后出現在數據集中的次數。可見高頻詞“the”的采樣率不足1/20。 ``` python def compare_counts(token): return '# %s: before=%d, after=%d' % (token, sum( [st.count(token_to_idx[token]) for st in dataset]), sum( [st.count(token_to_idx[token]) for st in subsampled_dataset])) compare_counts('the') # '# the: before=50770, after=2013' ``` 但低頻詞“join”則完整地保留了下來。 ``` python compare_counts('join') # '# join: before=45, after=45' ``` ### 10.3.1.3 提取中心詞和背景詞我們將與中心詞距離不超過背景窗口大小的詞作為它的背景詞。下面定義函數提取出所有中心詞和它們的背景詞。它每次在整數1和`max_window_size`（最大背景窗口）之間隨機均勻采樣一個整數作為背景窗口大小。 ``` python def get_centers_and_contexts(dataset, max_window_size): centers, contexts = [], [] for st in dataset: if len(st) < 2: # 每個句子至少要有2個詞才可能組成一對“中心詞-背景詞” continue centers += st for center_i in range(len(st)): window_size = random.randint(1, max_window_size) indices = list(range(max(0, center_i - window_size), min(len(st), center_i + 1 + window_size))) indices.remove(center_i) # 將中心詞排除在背景詞之外 contexts.append([st[idx] for idx in indices]) return centers, contexts ``` 下面我們創建一個人工數據集，其中含有詞數分別為7和3的兩個句子。設最大背景窗口為2，打印所有中心詞和它們的背景詞。 ``` python tiny_dataset = [list(range(7)), list(range(7, 10))] print('dataset', tiny_dataset) for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)): print('center', center, 'has contexts', context) ``` 輸出： ``` dataset [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9]] center 0 has contexts [1, 2] center 1 has contexts [0, 2, 3] center 2 has contexts [1, 3] center 3 has contexts [2, 4] center 4 has contexts [3, 5] center 5 has contexts [3, 4, 6] center 6 has contexts [4, 5] center 7 has contexts [8] center 8 has contexts [7, 9] center 9 has contexts [7, 8] ``` 實驗中，我們設最大背景窗口大小為5。下面提取數據集中所有的中心詞及其背景詞。 ``` python all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5) ``` ## 10.3.2 負采樣我們使用負采樣來進行近似訓練。對于一對中心詞和背景詞，我們隨機采樣`$ K $`個噪聲詞（實驗中設`$ K=5 $`）。根據word2vec論文的建議，噪聲詞采樣概率`$ P(w) $`設為`$ w $`詞頻與總詞頻之比的0.75次方 [2]。 ``` python 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: negatives = [] while len(negatives) < len(contexts) * K: 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): negatives.append(neg) all_negatives.append(negatives) return all_negatives sampling_weights = [counter[w]**0.75 for w in idx_to_token] all_negatives = get_negatives(all_contexts, sampling_weights, 5) ``` ## 10.3.3 讀取數據我們從數據集中提取所有中心詞`all_centers`，以及每個中心詞對應的背景詞`all_contexts`和噪聲詞`all_negatives`。我們先定義一個`Dataset`類。 ``` python class MyDataset(torch.utils.data.Dataset): def __init__(self, centers, contexts, negatives): assert len(centers) == len(contexts) == len(negatives) self.centers = centers self.contexts = contexts self.negatives = negatives def __getitem__(self, index): return (self.centers[index], self.contexts[index], self.negatives[index]) def __len__(self): return len(self.centers) ``` 我們將通過隨機小批量來讀取它們。在一個小批量數據中，第`$ 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。下面我們實現這個小批量讀取函數`batchify`。它的小批量輸入`data`是一個長度為批量大小的列表，其中每個元素分別包含中心詞`center`、背景詞`context`和噪聲詞`negative`。該函數返回的小批量數據符合我們需要的格式，例如，包含了掩碼變量。 ``` python def batchify(data): """用作DataLoader的參數collate_fn: 輸入是個長為batchsize的list, list中的每個元素都是Dataset類調用__getitem__得到的結果 """ max_len = max(len(c) + len(n) for _, c, n in data) centers, contexts_negatives, masks, labels = [], [], [], [] for center, context, negative in data: cur_len = len(context) + len(negative) centers += [center] 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))] return (torch.tensor(centers).view(-1, 1), torch.tensor(contexts_negatives), torch.tensor(masks), torch.tensor(labels)) ``` 我們用剛剛定義的`batchify`函數指定`DataLoader`實例中小批量的讀取方式，然后打印讀取的第一個批量中各個變量的形狀。 ``` python batch_size = 512 num_workers = 0 if sys.platform.startswith('win32') else 4 dataset = MyDataset(all_centers, all_contexts, all_negatives) data_iter = Data.DataLoader(dataset, batch_size, shuffle=True, collate_fn=batchify, num_workers=num_workers) for batch in data_iter: for name, data in zip(['centers', 'contexts_negatives', 'masks', 'labels'], batch): print(name, 'shape:', data.shape) break ``` 輸出： ``` centers shape: torch.Size([512, 1]) contexts_negatives shape: torch.Size([512, 60]) masks shape: torch.Size([512, 60]) labels shape: torch.Size([512, 60]) ``` ## 10.3.4 跳字模型我們將通過使用嵌入層和小批量乘法來實現跳字模型。它們也常常用于實現其他自然語言處理的應用。 ### 10.3.4.1 嵌入層獲取詞嵌入的層稱為嵌入層，在PyTorch中可以通過創建`nn.Embedding`實例得到。嵌入層的權重是一個矩陣，其行數為詞典大小（`num_embeddings`），列數為每個詞向量的維度（`embedding_dim`）。我們設詞典大小為20，詞向量的維度為4。 ``` python embed = nn.Embedding(num_embeddings=20, embedding_dim=4) embed.weight ``` 輸出： ``` Parameter containing: tensor([[-0.4689, 0.2420, 0.9826, -1.3280], [-0.6690, 1.2385, -1.7482, 0.2986], [ 0.1193, 0.1554, 0.5038, -0.3619], [-0.0347, -0.2806, 0.3854, -0.8600], [-0.6479, -1.1424, -1.1920, 0.3922], [ 0.6334, -0.0703, 0.0830, -0.4782], [ 0.1712, 0.8098, -1.2208, 0.4169], [-0.9925, 0.9383, -0.3808, -0.1242], [-0.3762, 1.9276, 0.6279, -0.6391], [-0.8518, 2.0105, 1.8484, -0.5646], [-1.0699, -1.0822, -0.6945, -0.7321], [ 0.4806, -0.5945, 1.0795, 0.1062], [-1.5377, 1.0420, 0.4325, 0.1098], [-0.8438, -1.4104, -0.9700, -0.4889], [-1.9745, -0.3092, 0.6398, -0.4368], [ 0.0484, -0.8516, -0.4955, -0.1363], [-2.6301, -0.7091, 2.2116, -0.1363], [-0.2025, 0.8037, 0.4906, 1.5929], [-0.6745, -0.8791, -0.9220, -0.8125], [ 0.2450, 1.9456, 0.1257, -0.3728]], requires_grad=True) ``` 嵌入層的輸入為詞的索引。輸入一個詞的索引`$ i $`，嵌入層返回權重矩陣的第`$ i $`行作為它的詞向量。下面我們將形狀為(2, 3)的索引輸入進嵌入層，由于詞向量的維度為4，我們得到形狀為(2, 3, 4)的詞向量。 ``` python x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long) embed(x) ``` 輸出： ``` tensor([[[-0.6690, 1.2385, -1.7482, 0.2986], [ 0.1193, 0.1554, 0.5038, -0.3619], [-0.0347, -0.2806, 0.3854, -0.8600]], [[-0.6479, -1.1424, -1.1920, 0.3922], [ 0.6334, -0.0703, 0.0830, -0.4782], [ 0.1712, 0.8098, -1.2208, 0.4169]]], grad_fn=<EmbeddingBackward>) ``` ### 10.3.4.2 小批量乘法我們可以使用小批量乘法運算`bmm`對兩個小批量中的矩陣一一做乘法。假設第一個小批量中包含`$ n $`個形狀為`$ a\times b $`的矩陣`$ \boldsymbol{X}_1, \ldots, \boldsymbol{X}_n $`，第二個小批量中包含`$ n $`個形狀為`$ b\times c $`的矩陣`$ \boldsymbol{Y}_1, \ldots, \boldsymbol{Y}_n $`。這兩個小批量的矩陣乘法輸出為`$ n $` 個形狀為`$ a\times c $`的矩陣`$ \boldsymbol{X}_1\boldsymbol{Y}_1, \ldots, \boldsymbol{X}_n\boldsymbol{Y}_n $`。因此，給定兩個形狀分別為(`$ n $`, `$ a $`, `$ b $`)和(`$ n $`, `$ b $`,`$ c $`)的`Tensor`，小批量乘法輸出的形狀為(`$ n $`, `$ a $`, `$ c $`)。 ``` python X = torch.ones((2, 1, 4)) Y = torch.ones((2, 4, 6)) torch.bmm(X, Y).shape ``` 輸出： ``` torch.Size([2, 1, 6]) ``` ### 10.3.4.3 跳字模型前向計算在前向計算中，跳字模型的輸入包含中心詞索引`center`以及連結的背景詞與噪聲詞索引`contexts_and_negatives`。其中`center`變量的形狀為(批量大小, 1)，而`contexts_and_negatives`變量的形狀為(批量大小, `max_len`)。這兩個變量先通過詞嵌入層分別由詞索引變換為詞向量，再通過小批量乘法得到形狀為(批量大小, 1, `max_len`)的輸出。輸出中的每個元素是中心詞向量與背景詞向量或噪聲詞向量的內積。 ``` python def skip_gram(center, contexts_and_negatives, embed_v, embed_u): v = embed_v(center) u = embed_u(contexts_and_negatives) pred = torch.bmm(v, u.permute(0, 2, 1)) return pred ``` ## 10.3.5 訓練模型在訓練詞嵌入模型之前，我們需要定義模型的損失函數。 ### 10.3.5.1 二元交叉熵損失函數根據負采樣中損失函數的定義，我們可以使用二元交叉熵損失函數,下面定義`SigmoidBinaryCrossEntropyLoss`。 ``` python class SigmoidBinaryCrossEntropyLoss(nn.Module): def __init__(self): # none mean sum super(SigmoidBinaryCrossEntropyLoss, self).__init__() def forward(self, inputs, targets, mask=None): """ input – Tensor shape: (batch_size, len) target – Tensor of the same shape as input """ inputs, targets, mask = inputs.float(), targets.float(), mask.float() res = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction="none", weight=mask) return res.mean(dim=1) loss = SigmoidBinaryCrossEntropyLoss() ``` 值得一提的是，我們可以通過掩碼變量指定小批量中參與損失函數計算的部分預測值和標簽：當掩碼為1時，相應位置的預測值和標簽將參與損失函數的計算；當掩碼為0時，相應位置的預測值和標簽則不參與損失函數的計算。我們之前提到，掩碼變量可用于避免填充項對損失函數計算的影響。 ``` python pred = torch.tensor([[1.5, 0.3, -1, 2], [1.1, -0.6, 2.2, 0.4]]) # 標簽變量label中的1和0分別代表背景詞和噪聲詞 label = torch.tensor([[1, 0, 0, 0], [1, 1, 0, 0]]) mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 0]]) # 掩碼變量 loss(pred, label, mask) * mask.shape[1] / mask.float().sum(dim=1) ``` 輸出： ``` tensor([0.8740, 1.2100]) ``` 作為比較，下面將從零開始實現二元交叉熵損失函數的計算，并根據掩碼變量`mask`計算掩碼為1的預測值和標簽的損失。 ``` python def sigmd(x): return - math.log(1 / (1 + math.exp(-x))) print('%.4f' % ((sigmd(1.5) + sigmd(-0.3) + sigmd(1) + sigmd(-2)) / 4)) # 注意1-sigmoid(x) = sigmoid(-x) print('%.4f' % ((sigmd(1.1) + sigmd(-0.6) + sigmd(-2.2)) / 3)) ``` 輸出： ``` 0.8740 1.2100 ``` ### 10.3.5.2 初始化模型參數我們分別構造中心詞和背景詞的嵌入層，并將超參數詞向量維度`embed_size`設置成100。 ``` python embed_size = 100 net = nn.Sequential( nn.Embedding(num_embeddings=len(idx_to_token), embedding_dim=embed_size), nn.Embedding(num_embeddings=len(idx_to_token), embedding_dim=embed_size) ) ``` ### 10.3.5.3 定義訓練函數下面定義訓練函數。由于填充項的存在，與之前的訓練函數相比，損失函數的計算稍有不同。 ``` python def train(net, lr, num_epochs): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print("train on", device) net = net.to(device) optimizer = torch.optim.Adam(net.parameters(), lr=lr) for epoch in range(num_epochs): start, l_sum, n = time.time(), 0.0, 0 for batch in data_iter: center, context_negative, mask, label = [d.to(device) for d in batch] pred = skip_gram(center, context_negative, net[0], net[1]) # 使用掩碼變量mask來避免填充項對損失函數計算的影響 l = (loss(pred.view(label.shape), label, mask) * mask.shape[1] / mask.float().sum(dim=1)).mean() # 一個batch的平均loss optimizer.zero_grad() l.backward() optimizer.step() l_sum += l.cpu().item() n += 1 print('epoch %d, loss %.2f, time %.2fs' % (epoch + 1, l_sum / n, time.time() - start)) ``` 現在我們就可以使用負采樣訓練跳字模型了。 ``` python train(net, 0.01, 10) ``` 輸出： ``` train on cpu epoch 1, loss 1.97, time 74.53s epoch 2, loss 0.62, time 81.85s epoch 3, loss 0.45, time 74.49s epoch 4, loss 0.39, time 72.04s epoch 5, loss 0.37, time 72.21s epoch 6, loss 0.35, time 71.81s epoch 7, loss 0.34, time 72.00s epoch 8, loss 0.33, time 74.45s epoch 9, loss 0.32, time 72.08s epoch 10, loss 0.32, time 72.05s ``` ## 10.3.6 應用詞嵌入模型訓練好詞嵌入模型之后，我們可以根據兩個詞向量的余弦相似度表示詞與詞之間在語義上的相似度。可以看到，使用訓練得到的詞嵌入模型時，與詞“chip”語義最接近的詞大多與芯片有關。 ``` python def get_similar_tokens(query_token, k, embed): W = embed.weight.data x = W[token_to_idx[query_token]] # 添加的1e-9是為了數值穩定性 cos = torch.matmul(W, x) / (torch.sum(W * W, dim=1) * torch.sum(x * x) + 1e-9).sqrt() _, topk = torch.topk(cos, k=k+1) topk = topk.cpu().numpy() for i in topk[1:]: # 除去輸入詞 print('cosine sim=%.3f: %s' % (cos[i], (idx_to_token[i]))) get_similar_tokens('chip', 3, net[0]) ``` 輸出： ``` cosine sim=0.478: hard-disk cosine sim=0.446: intel cosine sim=0.440: drives ``` ## 小結 * 可以使用PyTorch通過負采樣訓練跳字模型。 * 二次采樣試圖盡可能減輕高頻詞對訓練詞嵌入模型的影響。 * 可以將長度不同的樣本填充至長度相同的小批量，并通過掩碼變量區分非填充和填充，然后只令非填充參與損失函數的計算。 ## 參考文獻 [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). ----------- > 注：本節除代碼外與原書基本相同，[原書傳送門](https://zh.d2l.ai/chapter_natural-language-processing/word2vec-gluon.html)