预训练word2vec¶
:label:sec_word2vec_pretraining
我们继续实现
:numref:sec_word2vec
中定义的跳元语法模型。然后,我们将在PTB数据集上使用负采样预训练word2vec。首先,让我们通过调用d2l.load_data_ptb
函数来获得该数据集的数据迭代器和词表,该函数在
:numref:sec_word2vec_data
中进行了描述。
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)
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)
嵌入层的输入是词元(词)的索引。对于任何词元索引$i$,其向量表示可以从嵌入层中的权重矩阵的第$i$行获得。由于向量维度(output_dim
)被设置为4,因此当小批量词元索引的形状为(2,3)时,嵌入层返回具有形状(2,3,4)的向量。
x = torch.tensor([\[1, 2, 3], [4, 5, 6]\])
embed(x)
tensor([\[\[-1.4754, -0.3612, -0.4246, 0.5805], [-0.3160, 0.8830, 0.5328, 0.2179], [-0.0378, -0.5559, 1.4525, 0.6230]\], [\[ 0.0829, -1.0549, 0.6381, 0.7886], [-0.3862, -0.1291, 0.4160, -0.6710], [-0.4056, 0.0370, -0.6308, -0.2865]\]\], grad_fn=<EmbeddingBackward0>)
定义前向传播¶
在前向传播中,跳元语法模型的输入包括形状为(批量大小,1)的中心词索引center
和形状为(批量大小,max_len
)的上下文与噪声词索引contexts_and_negatives
,其中max_len
在
:numref:subsec_word2vec-minibatch-loading
中定义。这两个变量首先通过嵌入层从词元索引转换成向量,然后它们的批量矩阵相乘(在
:numref:subsec_batch_dot
中描述)返回形状为(批量大小,1,max_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
让我们为一些样例输入打印此skip_gram
函数的输出形状。
skip_gram(torch.ones((2, 1), dtype=torch.long),
torch.ones((2, 4), dtype=torch.long), embed, embed).shape
torch.Size([2, 1, 4])
class SigmoidBCELoss(nn.Module):
# 带掩码的二元交叉熵损失
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()
回想一下我们在 :numref:subsec_word2vec-minibatch-loading
中对掩码变量和标签变量的描述。下面计算给定变量的二进制交叉熵损失。
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])
下面显示了如何使用二元交叉熵损失中的Sigmoid激活函数(以较低效率的方式)计算上述结果。我们可以将这两个输出视为两个规范化的损失,在非掩码预测上进行平均。
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
初始化模型参数¶
我们定义了两个嵌入层,将词表中的所有单词分别作为中心词和上下文词使用。字向量维度embed_size
被设置为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))
定义训练阶段代码¶
训练阶段代码实现定义如下。由于填充的存在,损失函数的计算与以前的训练函数略有不同。
def train(net, data_iter, lr, num_epochs, device=d2l.try_gpu()):
def init_weights(m):
if type(m) == nn.Embedding:
nn.init.xavier_uniform_(m.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])
# 规范化的损失之和,规范化的损失数
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)}')
现在,我们可以使用负采样来训练跳元模型。
lr, num_epochs = 0.002, 5
train(net, data_iter, lr, num_epochs)
loss 0.410, 377799.5 tokens/sec on cuda:0
def get_similar_tokens(query_token, k, embed):
W = embed.weight.data
x = W[vocab[query_token]\]
# 计算余弦相似性。增加1e-9以获得数值稳定性
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:]: # 删除输入词
print(f'cosine sim={float(cos[i]):.3f}: {vocab.to_tokens(i)}')
get_similar_tokens('chip', 3, net[0])
cosine sim=0.773: microprocessor cosine sim=0.589: hitachi cosine sim=0.582: computers