预训练word2vec--Word2Vec实现(二)

预训练word2vec

现在，我们将在PTB数据集上使用负采样预训练word2vec。

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)

跳元模型

首先，回顾一下嵌入层是如何工作的。

嵌入层

如之前 所述，嵌入层将词元的索引映射到其特征向量。该层的权重是一个矩阵，其行数等于字典大小（input_dim）， 列数等于每个标记的向量维数（output_dim）。在词嵌入模型训练之后，这个权重就是我们所需要的。

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)

由于向量维度（output_dim）被设置为4，因此当小批量词元索引的形状为（2，3）时，嵌入层返回具有形状（2，3，4）的向量。

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

(torch.Size([2, 3]),
 torch.Size([2, 3, 4]),
 tensor([[[ 4.7132e-01,  4.9846e-01,  1.1359e+00, -5.7184e-01],
          [ 2.6345e-01,  8.8413e-01, -4.7359e-01, -9.0366e-01],
          [ 4.2658e-01, -6.1649e-01, -1.3205e+00, -7.0420e-01]],
 
         [[ 1.1838e+00, -1.4929e-01,  6.5329e-01,  6.1452e-01],
          [ 7.1647e-01, -2.7112e-01,  3.1809e-04,  1.9801e+00],
          [-2.9196e-01, -9.7819e-01,  2.8638e+00,  2.1714e-01]]],
        grad_fn=))

定义前向传播

在前向传播中，跳元语法模型的输入包括中心词索引center，其形状为（批量大小，1）。
上下文与噪声词索引contexts_and_negatives, 其形状为（批量大小，max_len），包括中心词和上下文词的嵌入矩阵 embed。

最终的返回形状为 （批量大小，1，max_len）。

def skip_gram(center, context_and_negatives, embed_v, embed_u):
    
    v = embed_v(center)
    u = embed_u(context_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()

回想一下之前对掩码变量与标签变量的描述。下面计算给定变量的二进制交叉熵损失。

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激活函数（以较低效率的方式） 计算上述结果。我们可以将这两个输出视为两个规范化的损失，在非掩码预测上进行平均。

# 对于该函数，x的值越大其对应的损失就越小
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)                               # 存放在CPU 或者 GPU 中
    
    # 定义优化算法, Adam优化算法
    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上下文标记
            center, context_negative, mask, label = [
                data.to(device) for data in batch]
            
            # 其最终的返回形状为（批量大小，1，max_len）,代表上下文词与中心词的点积结果
            # 结果越大，则其损失越小，两个词越相似
            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, 49261.0 tokens/sec on cpu

应用词嵌入

在训练word2vec模型之后，我们可以使用训练好模型中词向量的余弦相似度来从词表中找到与输入单词语义最相似的单词。

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)
    
    # 获取 前k+1 个最大值的索引值
    topk = torch.topk(cos, k=k+1)[1].cpu().numpy().astype('int32')
    
    # 根据索引值输出它们的余弦值和对应的token
    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.752:microprocessor
cosine sim=0.702:intel
cosine sim=0.626:chips

小结

1、我们可以使用嵌入层和二元交叉熵损失来训练带负采样的跳元模型。

2、词嵌入的应用包括基于词向量的余弦相似度为给定词找到语义相似的词。