详细的 Word2Vec 模型搭建笔记（pytorch版本）

前言

• 在看完沐神的 Dive-into-Deep-Learning 中的 Word2vec 之后自己也尝试动手实践了一遍，把一些不太好理解的代码转换成比较好理解的代码(虽然损失了一些速度)，下面是结合自己的理解的一次记录。

• 这里粘上 pytorch 的版本的 Dive 的来源
  https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter10_natural-language-processing/10.3_word2vec-pytorch

• 这里再粘上看到的一篇非常好的讲解Word2vec原理的CSDN文章

  https://blog.csdn.net/han_xiaoyang/article/details/89082129

0.准备工作

• 这个程序里面的变量数量比较多，因此这里先进行一个说明

  • counter 是一个字典，键表示单词，值表示单词出现的次数 (只记录了比出现5次更大的单词)。
  • word_to_idx 是一个字典，键表示单词，值表示单词的索引。
  • idx_to_word 是一个字典，键表示单词的索引，值表示单词。
  • submember_list_idx 是一个列表，表示二次采样后的各句子的索引。
  • centers_all 是一个列表，表示中心词。
  • contexts_all 是一个列表，表示每一个中心词的相邻词。
  • negatives_all 是一个列表，表示负采样后的列表。

• 导入对应的函数库

  import collections
  import math
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import numpy as np
  import time
  import random
  

1.首先进行数据的读取

• 数据来源：采样自华尔街日报文章的一个小数据集。
  Penn Tree Bank. https://catalog.ldc.upenn.edu/LDC99T42
• 数据形式
  
  这个数据集中句尾符为 < eos >，生僻词全用 < unk > 表示，数字则被替换成了 "N"。
• 根据这种数据的形式，读取时按行读取，每一行以单词为最小单位分隔开成为列表，将转换后的每一行（列表）再组成一个列表。

  # 这里将数据集按行读入，每一行按空格分开
  path = "/content/drive/My Drive/深度学习/ptb.train.txt"
  with open(path, "r") as read:
      readrow = read.readlines()
      readlist = [line.split() for line in readrow]
  

2.计算每个单词的数量

• collection.Counter 就是一个计数器，用来计数列表中每个单词出现了多少次，返回一个类型，需要使用 dict 转化成字典类型。
• 由于 collection.Counter 只能计数列表中各元素的个数，因此第一步要生成一个member_list 列表将全部的单词展开。
• 最后对单词进行筛选，只取出出现次数在 5 次及以上的单词，生成 counter 键为单词，值为该单词出现的次数。

  # 这里对读入的单词的数量进行统计
  # 这里将所有元素展成一个列表，用来进行下面 Counter 的计算
  member_list = [member for group in readlist for member in group]
  counter_row = dict(collections.Counter(member_list))
  # 下面将出现频率比较高的单词筛选出来，筛选标准频率 >=5
  counter = {}
  for member,times in counter_row.items():
      if times >= 5:
          counter[member] = times
  

3.建立一个单词的双向索引

• 一方面将单词索引到序号，另一方面将序号索引到单词。

  # 下面对筛选出来的单词建立一个索引(双向索引，一方面将单词索引到序号，另一方面将序号索引到单词)
  word_to_idx = dict([(member,idx) for idx, member in enumerate(counter.keys())])
  idx_to_word = dict([(word_to_idx[member],member) for member in word_to_idx.keys()])
  # 进行整体句子索引的替换(注意次数过少的就直接删去)
  member_list_idx = []
  for item in readlist:
      newlinelist = []
      for i in item:
          if i in counter.keys(): # 这里直接删去出现次数小于 5 次的单词
              newlinelist.append(word_to_idx[i])
      member_list_idx.append(newlinelist)
  

4.进行二次采样

• 进行二次采样的目的是为了减少那些出现次数过多的单词比如 ‘a’,‘the’ 等没有啥实际意义的单词的出现频率。
  删除单词$w_i$的概率为：$$P(w_i)=max(1-\sqrt{\frac{t}{f(w_i)}},0)$$
    ~  

  其中$f(w_i)$表示单词$w_i$出现的次数与全部单词出现的次数的比值： $$f(w_i)=\frac{n(w_i)}{{\sum }_{j}n(w_j)}$$
  很明显可以看出，当某一个单词比如 ‘a’ 出现的频率比较高的时候，$f(w_i)$ 比较大则$P(w_i)$ 也相应的比较大，因此这个单词更容易被删去。

• 这里采到一个坑：np.random.choices 的效率很低，比 random.choices 要慢不少，但是即使使用 random.choices 也要比书中采用的方法慢一些。

  #这里进行二次采样，将出现频率过高的单词（比如 a、the什么的）删去一些
  def discard(number):
      f =  counter[idx_to_word[number]]/ total_number #计算f(wi)
      p = max(1 - math.sqrt(t/f), 0) #计算p(wi)
      temp = int(random.choices([0,1], weights = [p,1.0-p])[0]) # 进行随机取样
      # False 表示留下这个单词，True 表示这个单词需要被删除
      if temp == 1:
          return False
      else:
          return True 
  
  # 这个是dive中的表达方式，和自己写的相比速度要快很多，感觉是choice的问题
  # def discard(idx):
  #     return random.uniform(0, 1) < 1 - math.sqrt(
  #         1e-4 / counter[idx_to_word[idx]] * total_number)
  
  t = 1e-4
  start_time = time.time()
  submember_list_idx = [] #  经过筛选的每句话的词向量
  total_number = sum(counter.values()) # 这里直接计算出全部单词的个数，减少后面的复杂度
  for item in member_list_idx:
      newlinelist = []
      for i in item:
          if discard(i) == False:
              newlinelist.append(i)
      submember_list_idx.append(newlinelist)
  end_time = time.time()
  print("进行二次采样花费的时间为：{:.3f}".format(end_time-start_time))
  
  max_window_size = 5
  centers_all, contexts_all = get_centers_and_context(submember_list_idx, max_window_size)
  print("中心词的个数为 {}".format(len(centers_all)))
  

5.进行中心词和背景词的提取

• 我们将与中心词距离不超过背景窗口大小的词作为它的背景词。下面定义函数提取出所有中心词和它们的背景词。它每次在整数1和 max_window_size（最大背景窗口）之间随机均匀采样一个整数作为背景窗口大小，注意这里的 max_window_size 是单边的长度，而不是双边的长度。
• 采用的提取方式是滑动框，每滑动一次就把中间的单词先取出来，然后把框内除去中间词的剩下单词作为背景词。
• 注意最后返回的 centers_all 是一层列表，而 contexts_all 是一个两层的列表。

  # 这里进行中心词和背景词的提取
  def get_centers_and_context(dataset, max_window_size):
      centers, context = [],[]
      for st in dataset:
          if len(st) < 2:
              continue 
          for center_id in range(len(st)):
              # 注意这里的 window_size 是单边的长度，而不是双边的长度
              window_size = random.randint(1,max_window_size)# 每次在1到最大窗口之间随机选择窗口长度进行采样
              low_bound = max(center_id - window_size, 0) # 采样下界
              high_bound = min(center_id + window_size + 1, len(st)) # 采样上界
              centers_real = st[center_id]
              context_slide = st[low_bound:high_bound] # 取出背景词和中心词
              context_slide.remove(centers_real) # 把中心词删去，只留下背景词
              centers.append(centers_real)
              context.append(context_slide)
      return centers, context
  
  max_window_size = 5
  centers_all, contexts_all = get_centers_and_context(submember_list_idx, max_window_size)
  

5.进行负采样

• 这里使用负采样进行近似训练，至于为什么要采用负采样可以参考上面给出的 csdn 博客中的文章，噪声词采样概率 $P(\omega )$ 设为 $\omega$ 词频与总词频之比的 $0.75$ 次方。

  # 这里进行负采样，K值表示针对每一个的负采样的个数
  def get_negatives(contexts_all, sampling_weights, K):
      negatives_all = []
      i = 0
      population = list(range(len(sampling_weights)))
      # 根据每个词的权重（sampling_weights）随机生成 k 个词的索引作为噪声词。
      choose_list = random.choices(population, weights = sampling_weights, k = int(1e5))
      for context in contexts_all:
          negative = []
          while (len(negative) < len(context) * K):
              neg = choose_list[i]
              i = i + 1
              if (i == len(choose_list)):
                  choose_list = random.choices(population, weights = sampling_weights, k = int(1e5)) # 如果噪声词用完了就重新生成
                  i = 0
              # 噪声词不能是背景词
              if (neg not in set(context)):
                  negative.append(neg)
          negatives_all.append(negative)
      return negatives_all
  
  sampling_weights = [counter[idx_to_word[i]]**0.75 for i in range(len(idx_to_word))]
  negatives_all = get_negatives(contexts_all, sampling_weights, 5)
  

6.完成数据的读取

• 为了方便后面数据的读取，这里先定义一个 Dataset 类，用来一次性把 centers，contexts，negatives 读取出来。

  # 这里定义一个 Dataset 类
  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$ 的长度都是不一样的，为了保证最后拼接长度是相同的，这里先求出 $\max \{m_i+n_i\}$，然后把所有的拼接向量都补齐到这个长度。
  • 构造掩码变量：为了指明每个拼接变量的真实长度，这里需要构造掩码变量 masks，masks 保证在真实长度位置是 1，其余补零位置是 0。
  • 构造标签：为了进一步分清背景词和噪声词，这里构造标签变量 labels，保证背景词处为 1，其余（包括噪声词和补齐的零）处为 0。

    # 这里定义了数据的读取方法
    def batchify(data):
        # 最后保证所有的用来训练的词向量的长度均相同，均为 max_len
        max_len = max(len(i) + len(j) for _,i,j in data) # 求出所有组合对中的长度最大值
        centers , labels, masks, combine_context = [], [], [], []
        # 每次都从数据中读取中心词，背景词和噪声词
        for center, context, negative in data:
            data_len = len(context) + len(negative) # 将背景词和噪声词组合成一个向量
            centers += [center] 
            combine_context += [context + negative + [0] * (max_len - data_len)]
            labels += [[1] * len(context) + [0] * (max_len - len(context))]
            masks += [[1] * data_len + [0] * (max_len - data_len)] 
        return (torch.tensor(centers).view(-1, 1), torch.tensor(combine_context), 
                torch.tensor(masks), torch.tensor(labels))
    

• 生成小批量的迭代器，选择批大小为 512，最后第一批的 max_len = 60。

  # 这里完成数据的批量化迭代生成,并且打印出生成的样本的情况
  batch_size = 512
  dataset = MyDataset(centers_all, contexts_all, negatives_all)
  dataiter = torch.utils.data.DataLoader(dataset, batch_size, True, collate_fn = batchify)
  for batch in dataiter:
      for name, data in zip(['centers', 'contexts_negatives', 'masks', 'labels'], batch):
          print(name, 'shape:', data.shape)
      break
  

  

7.定义 skip-gram 模型

• 编写 skip-gram 函数

  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(centers, combine_context, net[0], net[1])，定义的参数

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

综合函数和调用方式来看，这两个变量先通过词嵌入层分别由词索引变换为词向量，再通过小批量乘法得到形状为(批量大小, 1, max_len)的输出。输出中的每个元素是中心词向量与背景词向量或噪声词向量的内积。

8.定义二元交叉熵模型

• 具体为什么要用交叉熵可以参考 一文搞懂熵(Entropy),交叉熵(Cross-Entropy)，简单来说交叉熵对比了模型的预测结果和数据的真实标签，随着预测越来越准确，交叉熵的值越来越小，如果预测完全正确，交叉熵的值就为0。
• 在 torch.nn.functional 库中采用 binary_cross_entropy_with_logits来完成预测结果与真实标签之间交叉熵的计算，注意这个函数同时还具有 sigmoid 的功能所以后面不用额外添加归一化映射了。

  class SigmoidBinaryCrossEntropyLoss(nn.Module):
      def __init__(self):
          super(SigmoidBinaryCrossEntropyLoss, self).__init__()
      def forward(self, inputs, targets, mask = None):
          inputs, targets, mask = inputs.float(), targets.float(), mask.float()
          res = F.binary_cross_entropy_with_logits(inputs, targets, reduction = "none", weight = mask)
          return res.mean(dim=1)
  

9.完成训练

epochs = 10
lr = 0.01
optimizer = torch.optim.Adam(net.parameters(), lr = lr)
loss = SigmoidBinaryCrossEntropyLoss()
for epoch in range(epochs):
    l_sum = 0.0
    n = 0
    start = time.time()
    for batch in dataiter:
        centers, combine_context, masks, labels = batch
        pred = skip_gram(centers, combine_context, net[0], net[1])
        # 使用掩码变量mask来避免填充项对损失函数计算的影响
        l = (loss(pred.view(labels.shape), labels, masks) * masks.shape[1] / masks.float().sum(dim=1)).mean()# 一个batch的平均loss
        optimizer.zero_grad()
        l.backward()
        optimizer.step()
        l_sum += l.item()
        n += 1
    print("第{}次epoch，loss为{}，time为{}".format(epoch + 1, l_sum/n, time.time()-start))

10.应用词嵌入模型

• item 表示需要查找与之相似的词，是一个字符串；k 表示列出最相似的前 k 个词；embed 表示那个训练好的词嵌入模型。

  def get_similar_tokens(item, k, embed):
      W = embed.weight.data
      x = W[word_to_idx[item]]
      # 添加的 1e-9 是为了数值稳定性，防止除法出现 0
      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)
      for i in topk[1:]:
          print("余弦相似量为{:.3f},相似的词为{}".format(cos[i].item(), idx_to_word[int(i)]))
  

• 比如来看一下和“chip” 这个词类似的词，get_similar_tokens("chip", 5, net[0])