PyTorch教程-15.4. 預(yù)訓(xùn)練word2vec

我們繼續(xù)實現(xiàn) 15.1 節(jié)中定義的 skip-gram 模型。然后我們將在 PTB 數(shù)據(jù)集上使用負(fù)采樣來預(yù)訓(xùn)練 word2vec。首先，讓我們通過調(diào)用函數(shù)來獲取數(shù)據(jù)迭代器和這個數(shù)據(jù)集的詞匯表 ，這在第 15.3 節(jié)d2l.load_data_ptb中有描述

import math
import torch
from torch import nn
from d2l import torch as d2l

batch_size, max_window_size, num_noise_words = 512, 5, 5
data_iter, vocab = d2l.load_data_ptb(batch_size, max_window_size,
                   num_noise_words)

Downloading ../data/ptb.zip from http://d2l-data.s3-accelerate.amazonaws.com/ptb.zip...

import math
from mxnet import autograd, gluon, np, npx
from mxnet.gluon import nn
from d2l import mxnet as d2l

npx.set_np()

batch_size, max_window_size, num_noise_words = 512, 5, 5
data_iter, vocab = d2l.load_data_ptb(batch_size, max_window_size,
                   num_noise_words)

15.4.1。Skip-Gram 模型

我們通過使用嵌入層和批量矩陣乘法來實現(xiàn) skip-gram 模型。首先，讓我們回顧一下嵌入層是如何工作的。

15.4.1.1。嵌入層

如第 10.7 節(jié)所述，嵌入層將標(biāo)記的索引映射到其特征向量。該層的權(quán)重是一個矩陣，其行數(shù)等于字典大小 ( input_dim)，列數(shù)等于每個標(biāo)記的向量維數(shù) ( output_dim)。一個詞嵌入模型訓(xùn)練好之后，這個權(quán)重就是我們所需要的。

embed = nn.Embedding(num_embeddings=20, embedding_dim=4)
print(f'Parameter embedding_weight ({embed.weight.shape}, '
   f'dtype={embed.weight.dtype})')

Parameter embedding_weight (torch.Size([20, 4]), dtype=torch.float32)

embed = nn.Embedding(input_dim=20, output_dim=4)
embed.initialize()
embed.weight

Parameter embedding0_weight (shape=(20, 4), dtype=float32)

嵌入層的輸入是標(biāo)記（單詞）的索引。對于任何令牌索引i，它的向量表示可以從ith嵌入層中權(quán)重矩陣的行。由于向量維度 ( output_dim) 設(shè)置為 4，因此嵌入層返回形狀為 (2, 3, 4) 的向量，用于形狀為 (2, 3) 的標(biāo)記索引的小批量。

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
embed(x)

tensor([[[-0.6501, 1.3547, 0.7968, 0.3916],
     [ 0.4739, -0.0944, 1.2308, 0.6457],
     [ 0.4539, 1.5194, 0.4377, -1.5122]],

    [[-0.7032, -0.1213, 0.2657, -0.6797],
     [ 0.2930, -0.6564, 0.8960, -0.5637],
     [-0.1815, 0.9487, 0.8482, 0.5486]]], grad_fn=)

x = np.array([[1, 2, 3], [4, 5, 6]])
embed(x)

array([[[ 0.01438687, 0.05011239, 0.00628365, 0.04861524],
    [-0.01068833, 0.01729892, 0.02042518, -0.01618656],
    [-0.00873779, -0.02834515, 0.05484822, -0.06206018]],

    [[ 0.06491279, -0.03182812, -0.01631819, -0.00312688],
    [ 0.0408415 , 0.04370362, 0.00404529, -0.0028032 ],
    [ 0.00952624, -0.01501013, 0.05958354, 0.04705103]]])

15.4.1.2。定義前向傳播

在正向傳播中，skip-gram 模型的輸入包括形狀為（批大小，1）的中心詞索引和 形狀為（批大小，）center的連接上下文和噪聲詞索引，其中定義在 第 15.3.5 節(jié). 這兩個變量首先通過嵌入層從標(biāo)記索引轉(zhuǎn)換為向量，然后它們的批量矩陣乘法（在第 11.3.2.2 節(jié)中描述）返回形狀為（批量大小，1， ）的輸出 。輸出中的每個元素都是中心詞向量與上下文或噪聲詞向量的點積。contexts_and_negativesmax_lenmax_lenmax_len

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

def skip_gram(center, contexts_and_negatives, embed_v, embed_u):
  v = embed_v(center)
  u = embed_u(contexts_and_negatives)
  pred = npx.batch_dot(v, u.swapaxes(1, 2))
  return pred

skip_gram讓我們?yōu)橐恍┦纠斎氪蛴〈撕瘮?shù)的輸出形狀。

skip_gram(torch.ones((2, 1), dtype=torch.long),
     torch.ones((2, 4), dtype=torch.long), embed, embed).shape

torch.Size([2, 1, 4])

skip_gram(np.ones((2, 1)), np.ones((2, 4)), embed, embed).shape

(2, 1, 4)

15.4.2。訓(xùn)練

在用負(fù)采樣訓(xùn)練skip-gram模型之前，我們先定義它的損失函數(shù)。

15.4.2.1。二元交叉熵?fù)p失

根據(jù)15.2.1節(jié)負(fù)采樣損失函數(shù)的定義，我們將使用二元交叉熵?fù)p失。

class SigmoidBCELoss(nn.Module):
  # Binary cross-entropy loss with masking
  def __init__(self):
    super().__init__()

  def forward(self, inputs, target, mask=None):
    out = nn.functional.binary_cross_entropy_with_logits(
      inputs, target, weight=mask, reduction="none")
    return out.mean(dim=1)

loss = SigmoidBCELoss()

loss = gluon.loss.SigmoidBCELoss()

回想我們在第 15.3.5 節(jié)中對掩碼變量和標(biāo)簽變量的描述 。下面計算給定變量的二元交叉熵?fù)p失。

pred = torch.tensor([[1.1, -2.2, 3.3, -4.4]] * 2)
label = torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]])
mask = torch.tensor([[1, 1, 1, 1], [1, 1, 0, 0]])
loss(pred, label, mask) * mask.shape[1] / mask.sum(axis=1)

tensor([0.9352, 1.8462])

pred = np.array([[1.1, -2.2, 3.3, -4.4]] * 2)
label = np.array([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]])
mask = np.array([[1, 1, 1, 1], [1, 1, 0, 0]])
loss(pred, label, mask) * mask.shape[1] / mask.sum(axis=1)

array([0.93521017, 1.8462094 ])

下面顯示了如何使用二元交叉熵?fù)p失中的 sigmoid 激活函數(shù)計算上述結(jié)果（以效率較低的方式）。我們可以將這兩個輸出視為兩個歸一化損失，對非掩碼預(yù)測進(jìn)行平均。

def sigmd(x):
  return -math.log(1 / (1 + math.exp(-x)))

print(f'{(sigmd(1.1) + sigmd(2.2) + sigmd(-3.3) + sigmd(4.4)) / 4:.4f}')
print(f'{(sigmd(-1.1) + sigmd(-2.2)) / 2:.4f}')

0.9352
1.8462

def sigmd(x):
  return -math.log(1 / (1 + math.exp(-x)))

print(f'{(sigmd(1.1) + sigmd(2.2) + sigmd(-3.3) + sigmd(4.4)) / 4:.4f}')
print(f'{(sigmd(-1.1) + sigmd(-2.2)) / 2:.4f}')

0.9352
1.8462

15.4.2.2。初始化模型參數(shù)

當(dāng)詞匯表中的所有詞分別用作中心詞和上下文詞時，我們?yōu)樗鼈兌x了兩個嵌入層。詞向量維度embed_size設(shè)置為100。

embed_size = 100
net = nn.Sequential(nn.Embedding(num_embeddings=len(vocab),
                 embedding_dim=embed_size),
          nn.Embedding(num_embeddings=len(vocab),
                 embedding_dim=embed_size))

embed_size = 100
net = nn.Sequential()
net.add(nn.Embedding(input_dim=len(vocab), output_dim=embed_size),
    nn.Embedding(input_dim=len(vocab), output_dim=embed_size))

15.4.2.3。定義訓(xùn)練循環(huán)

訓(xùn)練循環(huán)定義如下。由于padding的存在，損失函數(shù)的計算與之前的訓(xùn)練函數(shù)略有不同。

def train(net, data_iter, lr, num_epochs, device=d2l.try_gpu()):
  def init_weights(module):
    if type(module) == nn.Embedding:
      nn.init.xavier_uniform_(module.weight)
  net.apply(init_weights)
  net = net.to(device)
  optimizer = torch.optim.Adam(net.parameters(), lr=lr)
  animator = d2l.Animator(xlabel='epoch', ylabel='loss',
              xlim=[1, num_epochs])
  # Sum of normalized losses, no. of normalized losses
  metric = d2l.Accumulator(2)
  for epoch in range(num_epochs):
    timer, num_batches = d2l.Timer(), len(data_iter)
    for i, batch in enumerate(data_iter):
      optimizer.zero_grad()
      center, context_negative, mask, label = [
        data.to(device) for data in batch]

      pred = skip_gram(center, context_negative, net[0], net[1])
      l = (loss(pred.reshape(label.shape).float(), label.float(), mask)
           / mask.sum(axis=1) * mask.shape[1])
      l.sum().backward()
      optimizer.step()
      metric.add(l.sum(), l.numel())
      if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
        animator.add(epoch + (i + 1) / num_batches,
               (metric[0] / metric[1],))
  print(f'loss {metric[0] / metric[1]:.3f}, '
     f'{metric[1] / timer.stop():.1f} tokens/sec on {str(device)}')

def train(net, data_iter, lr, num_epochs, device=d2l.try_gpu()):
  net.initialize(ctx=device, force_reinit=True)
  trainer = gluon.Trainer(net.collect_params(), 'adam',
              {'learning_rate': lr})
  animator = d2l.Animator(xlabel='epoch', ylabel='loss',
              xlim=[1, num_epochs])
  # Sum of normalized losses, no. of normalized losses
  metric = d2l.Accumulator(2)
  for epoch in range(num_epochs):
    timer, num_batches = d2l.Timer(), len(data_iter)
    for i, batch in enumerate(data_iter):
      center, context_negative, mask, label = [
        data.as_in_ctx(device) for data in batch]
      with autograd.record():
        pred = skip_gram(center, context_negative, net[0], net[1])
        l = (loss(pred.reshape(label.shape), label, mask) *
           mask.shape[1] / mask.sum(axis=1))
      l.backward()
      trainer.step(batch_size)
      metric.add(l.sum(), l.size)
      if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
        animator.add(epoch + (i + 1) / num_batches,
               (metric[0] / metric[1],))
  print(f'loss {metric[0] / metric[1]:.3f}, '
     f'{metric[1] / timer.stop():.1f} tokens/sec on {str(device)}')

現(xiàn)在我們可以使用負(fù)采樣來訓(xùn)練 skip-gram 模型。

lr, num_epochs = 0.002, 5
train(net, data_iter, lr, num_epochs)

loss 0.410, 170397.2 tokens/sec on cuda:0

lr, num_epochs = 0.002, 5
train(net, data_iter, lr, num_epochs)

loss 0.408, 86715.1 tokens/sec on gpu(0)

15.4.3。應(yīng)用詞嵌入

在訓(xùn)練完 word2vec 模型后，我們可以使用來自訓(xùn)練模型的詞向量的余弦相似度來從詞典中找到與輸入詞在語義上最相似的詞。

def get_similar_tokens(query_token, k, embed):
  W = embed.weight.data
  x = W[vocab[query_token]]
  # Compute the cosine similarity. Add 1e-9 for numerical stability
  cos = torch.mv(W, x) / torch.sqrt(torch.sum(W * W, dim=1) *
                   torch.sum(x * x) + 1e-9)
  topk = torch.topk(cos, k=k+1)[1].cpu().numpy().astype('int32')
  for i in topk[1:]: # Remove the input words
    print(f'cosine sim={float(cos[i]):.3f}: {vocab.to_tokens(i)}')

get_similar_tokens('chip', 3, net[0])

cosine sim=0.665: microprocessor
cosine sim=0.665: chips
cosine sim=0.646: intel

def get_similar_tokens(query_token, k, embed):
  W = embed.weight.data()
  x = W[vocab[query_token]]
  # Compute the cosine similarity. Add 1e-9 for numerical stability
  cos = np.dot(W, x) / np.sqrt(np.sum(W * W, axis=1) * np.sum(x * x) + 1e-9)
  topk = npx.topk(cos, k=k+1, ret_typ='indices').asnumpy().astype('int32')
  for i in topk[1:]: # Remove the input words
    print(f'cosine sim={float(cos[i]):.3f}: {vocab.to_tokens(i)}')

get_similar_tokens('chip', 3, net[0])

cosine sim=0.689: intel
cosine sim=0.682: desktop
cosine sim=0.616: digital

15.4.4。概括

我們可以使用嵌入層和二元交叉熵?fù)p失來訓(xùn)練具有負(fù)采樣的 skip-gram 模型。

詞嵌入的應(yīng)用包括根據(jù)詞向量的余弦相似度為給定詞尋找語義相似的詞。

15.4.5。練習(xí)

使用經(jīng)過訓(xùn)練的模型，為其他輸入詞找到語義相似的詞。你能通過調(diào)整超參數(shù)來改善結(jié)果嗎？

當(dāng)訓(xùn)練語料庫很大時，我們在更新模型參數(shù)時，往往會在當(dāng)前minibatch中對中心詞進(jìn)行上下文詞和噪聲詞的采樣。換句話說，同一個中心詞在不同的訓(xùn)練時期可能有不同的上下文詞或噪聲詞。這種方法有什么好處？嘗試實施這種訓(xùn)練方法。