pytorch学习5-word2vec（skip-gram）实现

1.准备训练数据

（1）读取文本数据

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as tud
from torch.nn.parameter import Parameter

from collections import Counter
import numpy as np
import random
import math

import pandas as pd
import scipy
import sklearn
from sklearn.metrics.pairwise import cosine_similarity

USE_CUDA = torch.cuda.is_available()
# 为了保证实验结果可以复现，我们经常会把各种random seed固定在某一个值
random.seed(53113)
np.random.seed(53113)
torch.manual_seed(53113)
if USE_CUDA:
    torch.cuda.manual_seed(53113)
    
# 设定一些超参数
K = 100 # number of negative samples
C = 3 # nearby words threshold
NUM_EPOCHS = 2 # The number of epochs of training
MAX_VOCAB_SIZE = 30000 # the vocabulary size
BATCH_SIZE = 128 # the batch size
LEARNING_RATE = 0.2 # the initial learning rate
EMBEDDING_SIZE = 100

LOG_FILE = 'word-embedding.log'

# tokenize函数，把一篇文本转化成一个个单词
def word_tokenize(text):
    return text.split()
# 从文本文件中读取所有的文字，通过这些文本创建一个vocabulary
# 由于单词数量可能太大，我们只选取最常见的MAX_VOCAB_SIZE个单词
# 我们添加一个UNK单词表示所有不常见的单词
# 我们需要记录单词到index的mapping，以及index到单词的mapping，单词的count，单词的(normalized) frequency，以及单词总数。

with open('text8.train.txt','r') as fin:
	# 按照字节读取，返回一个字符串
    text = fin.read()
    
text = [w for w in word_tokenize(text.lower())]
# #建立词典，最后一位留给不常用或者没出现过的单词
vocab = dict(Counter(text).most_common(MAX_VOCAB_SIZE-1))
vocab['<unk>'] = len(text) - np.sum(list(vocab.values()))
 #建立词典，把key拿出来
idx_to_word = [word for word in vocab.keys()]
#字典{the:1, of:2, I:3, love:4, he:5, you:6}
word_to_idx = {word:i for i,word in enumerate(idx_to_word)}

word_counts = np.array([count for count in vocab.values()],dtype=np.float32)
word_freqs = word_counts/np.sum(word_counts)
word_freqs = word_freqs ** (3./4.)
word_freqs = word_freqs / np.sum(word_freqs)
VOCAB_SIZE = len(idx_to_word)
print(VOCAB_SIZE)

（2）实现dataloader

# 实现Dataloader¶
# 一个dataloader需要以下内容：

# 把所有text编码成数字，然后用subsampling预处理这些文字。
# 保存vocabulary，单词count，normalized word frequency
# 每个iteration sample一个中心词
# 根据当前的中心词返回context单词
# 根据中心词sample一些negative单词
# 返回单词的counts
# 有了dataloader之后，我们可以轻松随机打乱整个数据集，拿到一个batch的数据等等。
class WordEmbeddingDataset(tud.Dataset):
    def __init__(self,text,word_to_idx,idx_to_word,word_freqs,word_counts):
        ''' text: a list of words, all text from the training dataset
            word_to_idx: the dictionary from word to idx
            idx_to_word: idx to word mapping
            word_freq: the frequency of each word
            word_counts: the word counts
        '''
        super(WordEmbeddingDataset,self).__init__()
        # 把 text 文章中的所有词替换成 字典中的编号，get()函数，设置默认值
        self.text_encoded = [word_to_idx.get(t,VOCAB_SIZE-1) for t in text]
        # 把numpy 类型的变量转为 tensor类型
        self.text_encoded = torch.Tensor(self.text_encoded).long()
        self.word_to_idx = word_to_idx
        self.idx_to_word = idx_to_word
        #记录字典单词出现的频率 转为 tensor类型
        self.word_freqs = torch.Tensor(word_freqs)
        #记录字典单词出现的次数 转为 tensor类型
        self.word_counts = torch.Tensor(word_counts)
    def __len__(self):
        ''' 返回整个数据集（所有单词）的长度
        '''
        return len(self.text_encoded)
    def __getitem__(self,idx):
        ''' 这个function返回以下数据用于训练
            - 中心词
            - 这个单词附近的(positive)单词
            - 随机采样的K个单词作为negative sample
        '''
        #找到中心词在text中的位置 eq. love 为 4 .
        center_word = self.text_encoded[idx]
        # 找到在文中 中心词附近的 positive sample 这里C= 3. 有2*C 个 positive sample
        pos_indices = list(range(idx-C,idx)) + list(range(idx+1,idx+C+1))
        # 如果 中心词在开头或者结尾，那么我们把负数转到text的结尾或者是开头。eq -1%10 = 9
        pos_indices = [ i%len(self.text_encoded) for i in pos_indices]
        #找到中心词对应 positive 单词在字典中的编号
        pos_words = self.text_encoded[pos_indices]
        #负采样，根据权重采样，K*pos个单词,抽取的是word_freqs的下标索引值。
        neg_words = torch.multinomial(self.word_freqs, K * pos_words.shape[0], True)
        #把中心词， postive samples and negative samples 的编号返回
        return center_word,pos_words,neg_words

dataset = WordEmbeddingDataset(text,word_to_idx,idx_to_word,word_freqs,word_counts)
#num_workers 使用进程数，0 为1个进程。
dataloader = tud.DataLoader(dataset,batch_size=BATCH_SIZE,shuffle=True,num_workers=0)
for i,(input_labels,pos_labels,neg_labels) in enumerate(dataloader):
    print("------------------",i,"-----------------------")
    print(input_labels)
    print(pos_labels)
    print(neg_labels)

2.定义模型

# 定义PyTorch模型
class EmbeddingModel(nn.Module):
    def __init__(self,vocab_size,embed_size):
        ''' 初始化输出和输出embedding
        '''
        super(EmbeddingModel,self).__init__()
        #字典大小 30000
        self.vocab_size = vocab_size
        #单词维度 一般是50，100，300维
        self.embed_size = embed_size
        initrange = 0.5/self.embed_size
        #初始化一个矩阵，self.vocab_size * self.embed_size，sparse默认为false，如果为true，是稀疏矩阵，对于参数很大，部分为0会有用。
        self.out_embed = nn.Embedding(self.vocab_size,self.embed_size,sparse=False)
        #把矩阵中的权重初始化 设置在 -0.5 / self.embed_size到0.5 / self.embed_size间的随机值，经验值,会提升效率。
        self.out_embed.weight.data.uniform_(-initrange,initrange)
        #初始化一个矩阵，self.vocab_size * self.embed_size
        self.in_embed = nn.Embedding(self.vocab_size,self.embed_size,sparse=False)
        #把矩阵中的权重初始化 设置在 -0.5 / self.embed_size到0.5 / self.embed_size间的随机值
        self.in_embed.weight.data.uniform_(-initrange,initrange)
    
    def forward(self,input_labels,pos_labels,neg_labels):
        '''
            input_labels: 中心词, [batch_size,]
            pos_labels: 中心词周围 context window 出现过的单词 [batch_size , (window_size * 2)]
            neg_labelss: 中心词周围没有出现过的单词，从 negative sampling 得到 [batch_size, (window_size * 2 * K)]
            return: loss, [batch_size]
        '''
        batch_size = input_labels.size(0)
        # embedding方法获取的向量只是随机初始化的，并不代表任何含义，并且不会有word2vec等训练出来的效果。
        # 但是可以利用这样的方法先赋值然后在学习。
        # [batch_szie,embed_size]
        input_embedding = self.in_embed(input_labels)
        # [batch,(2*C),embed_size]
        pos_embedding = self.out_embed(pos_labels)
        # [batch, (2*C * K),embed_size]
        neg_embedding = self.out_embed(neg_labels)
        
        # [batch_size,(2*C)] torch,bmm()乘法，unsqueeze()增加维度，squeeze（）减少维度
        log_pos = torch.bmm(pos_embedding,input_embedding.unsqueeze(2)).squeeze()
        # [batch,(2*C*K)]
        log_neg = torch.bmm(neg_embedding, -input_embedding.unsqueeze(2)).squeeze()
        # torch.sum(1)，指定维度的tensor相加，维度会少一，log_pos变成一维向量，dim =1 按照行相加。
        log_pos = F.logsigmoid(log_pos).sum(1)
        log_neg = F.logsigmoid(log_neg).sum(1)
        # 对应位置相加
        loss = log_pos + log_neg
        return -loss
    def input_embeddings(self):
        return self.in_embed.weight.data.cpu().numpy()

3.定义评估函数

def evaluate(filename,embedding_weights):
    if filename.endswith(".csv"):
        data = pd.read_csv(filename,sep=',')
    else:
        data = pd.read_csv(filename,sep='\t')
    human_similarity = []
    model_similarity = []
    for i in data.iloc[:,0:2].index:
        word1,word2 = data.iloc[i,0],data.iloc[i,1]
        if word1 not in word_to_idx or word2 not in word_to_idx:
            continue
        else:
            word1_idx,word2_idx = word_to_idx[word1],word_to_idx[word2]
            word1_embed,word2_embed = embedding_weights[[word1_idx]],embedding_weights[[word2_idx]]
            #  计算余弦相似度            
            model_similarity.append(float(sklearn.metrics.pairwise.cosine_similarity(word1_embed,word2_embed)))
            human_similarity.append(float(data.iloc[i,2]))
    # spearman秩相关系数是度量两个变量之间的统计相关性的指标，用来评估当用单调函数来描述两个变量之间的关系有多好。
    return scipy.stats.spearmanr(human_similarity,model_similarity)

4.定义优化函数

model = EmbeddingModel(VOCAB_SIZE,EMBEDDING_SIZE)
if USE_CUDA:
    model = model.cuda()
optimizer = torch.optim.SGD(model.parameters(),lr=LEARNING_RATE)

5.训练模型

# 训练模型
for e in range(NUM_EPOCHS):
    t = 0
    for i,(input_labels,pos_labels,neg_labels) in enumerate(dataloader):
        input_labels = input_labels.long()
        pos_labels = pos_labels.long()
        neg_labels = neg_labels.long()
        if USE_CUDA:
            input_labels = input_labels.cuda()
            pos_labels = pos_labels.cuda()
            neg_labels = neg_labels.cuda()
        optimizer.zero_grad()
        loss = model(input_labels,pos_labels,neg_labels).mean()
        loss.backward()
        optimizer.step()
        
        if i % 100 == 0:
            with open(LOG_FILE,"a") as fout:
                fout.write("epoch:{},iter:{},loss:{}\n".format(e,i,loss.item()))
                print("epoch:{},iter:{},loss:{}".format(e,i,loss.item()))
        
        if i % 2000 == 0:
            embedding_weights = model.input_embeddings()
            sim_simlex = evaluate("simlex-999.txt",embedding_weights)
            sim_men = evaluate("men.txt",embedding_weights)
            sim_353 = evaluate("wordsim353.csv",embedding_weights)           
            with open(LOG_FILE,"a") as fout:
                print("epoch: {}, iteration: {}, simlex-999: {}, men: {}, sim353: {}, nearest to monster: {}\n".format(
                    e, i, sim_simlex, sim_men, sim_353, find_nearest("monster")))
                fout.write("epoch: {}, iteration: {}, simlex-999: {}, men: {}, sim353: {}, nearest to monster: {}\n".format(
                    e, i, sim_simlex, sim_men, sim_353, find_nearest("monster")))
    embedding_weights = model.input_embeddings()
    np.save("embedding--{}".format(EMBEDDING_SZIE),embedding_weights)
    torch.save(model.state_dict(),"embedding-{}.th".format(EMBEDDING_SIZE))

6.准确度评估

def find_nearest(word):
    index = word_to_idx[word]
    embedding = embedding_weights[index]
    cos_dis = np.array([scipy.spatial.distance.cosine(e,embedding) for e in embedding_weights])
    return [idx_to_word[i] for i in cos_dis.argsort()[:10]]
for word in ["good", "fresh", "monster", "green", "like", "america", "chicago", "work", "computer", "language"]:
    print(word, find_nearest(word))
man_idx = word_to_idx["man"] 
king_idx = word_to_idx["king"] 
woman_idx = word_to_idx["woman"]
embedding = embedding_weights[woman_idx] - embedding_weights[man_idx] + embedding_weights[king_idx]
cos_dis = np.array([scipy.spatial.distance.cosine(e, embedding) for e in embedding_weights])
for i in cos_dis.argsort()[:20]:
    print(idx_to_word[i])