Word2Vec的pytorch实现（Skip-gram）

写在前面：

本篇文章是我个人的学习记录，仅包含代码实现和一些个人理解，参考的一些文章我会给出链接。
深度学习word2vec笔记之算法篇.
Word2Vec
PyTorch 实现 Word2Vec

Word2Vec的数学原理详解：链接：https://pan.baidu.com/s/1_xlOvItHKYh4ndu4UaMF8Q
提取码：gt2f

相关代码：https://github.com/deborausujono/word2vecpy
如果文章出现错误或者有可优化的部分，都欢迎大家指出，一起学习进步~

正文

需要用的数据，文件中的内容是英文文本，去除了标点符号，每个单词之间用空格隔开
语料库下载地址：https://pan.baidu.com/s/10Bd3JxCCFTjBPNt0YROvZA
提取码：81fo

使用Skip-gram模型和负采样训练法。

代码部分

读取数据

首先定义一个类，读取原始语料，创建单词和id的映射关系，并计算词频。
这里词频使用了3/4 率.为 word2vec 论文里面推荐这么做

class Process_source_data():
    # 该类的作用是读取原始数据，构建词典映射关系和单词词频
    def __init__(self, data_path, vocab_size=10000):
        with open(data_path, 'r', encoding='utf8') as f:
            text = f.read()
        # 将单词切分
        self.text = text.lower().split()
        
        # 选出最多的vocab_size个词
        # 得到单词字典表，key是单词，value是次数
        vocob_dict = dict(Counter(self.text).most_common(vocab_size-1))
        # 把不常用的单词都编码为"<UNK>"
        vocob_dict['<UNK>'] = len(text) - np.sum(list(vocob_dict.values()))       
        
        #构建映射关系
        self.word2id = {word:i for i, word in enumerate(vocob_dict)}
        self.id2word = {i:word for i, word in enumerate(vocob_dict)}     
        
        # 根据3/4率调整词频
        word_counts = np.asarray(list(vocob_dict.values()))
        self.word_freq = (word_counts / np.sum(word_counts))** (3./4.)

构建DataSet

这里按照Skip-gram的方式构建训练数据，同时构建负采样样本。
继承的父类是torch.utils.data.Dateset，因此需要重新两个函数：def _len_ 和 def _getitem_

class Skip_gram_Dataset(Data.Dataset):
    # 该类的作用是将文字转换为对应ID， 并返回给定idx时对应的训练数据
    def __init__(self, text, word2id, word_freq, C=3, K=15, batch_size=32):
        super().__init__()
        self.C = C     # 上下文窗口
        self.K = K     # 负采样比例
        
        self.text_encoded = [word2id.get(word, word2id['<UNK>']) for word in text]
        self.text_encoded = torch.tensor(self.text_encoded, dtype=torch.long)
        
        self.word_freq = torch.tensor(word_freq)
        
    def __len__(self):
        return len(self.text_encoded)
    
    def __getitem__(self, idx):
        '''
        对于给定的idx，返回对应的训练数据
        - 中心词
        - 这个单词附近的positive word
        - 随机采样的K个单词作为negative word
        '''
        center_word = self.text_encoded[idx]
        pos_idx = list(range(idx - self.C, idx)) + list(range(idx + 1, idx + self.C + 1))
        pos_idx = [i % len(self.text_encoded) for i in pos_idx]
        pos_words = self.text_encoded[pos_idx]
        
        # torch.multinomial(input, num_samples, replacement=False, *, generator=None, out=None)
        # 可以根据input的权重随机选出num_samples个input的下标， 当input中的值为0时，则不会返回该值的下标
        # replacement表示是否是有放回的抽取
        select_weight = copy.deepcopy(self.word_freq)
        select_weight[pos_words] = 0    # 去除背景词
        select_weight[center_word] = 0   # 去除中心词
        # 每取一个背景词，需要取K倍的负采样
        neg_words = torch.multinomial(select_weight, self.K * pos_words.shape[0], True)

        return center_word, pos_words, neg_words  

构建Loader

利用torch.utils.data.DataLoader将dataset的数据按照batch_size输出。
这里将前两个类添加进来，实现从读取文本到构建数据迭代器的一条龙

class My_data_loader():
    def __init__(self, data_path, batch_size, shuffle=True):
        process_data = Process_source_data(data_path)
        self.dataset = Skip_gram_Dataset(process_data.text, process_data.word2id, process_data.word_freq)
        self.loader = Data.DataLoader(self.dataset, batch_size, shuffle)

构建训练模型

由于一个词可能充当中心词，也可能充当背景词，因此需要构建两个Embedding表分别进行词嵌入。
模型中的具体计算方法可查看文章开头给的两篇文章。
按照 Word2Vec 论文所写，推荐使用中心词向量作为最终的Embedding结果返回。

class Embedding_Model(nn.Module):
    def __init__(self, vocab_size, embed_size):
        super().__init__()
        self.in_embed = nn.Embedding(vocab_size, embed_size)  #中心词的词向量矩阵
        self.out_embed = nn.Embedding(vocab_size, embed_size) #背景词的词向量矩阵
        
    def forward(self, input_labels, pos_labels, neg_labels):
        input_embedding = self.in_embed(input_labels)  # [bs, embed_size]
        input_embedding = input_embedding.unsqueeze(2) # [bs, embed_size, 1]
        pos_embedding = self.out_embed(pos_labels)    # [bs, windows * 2 , embed_size]
        neg_embedding = self.out_embed(neg_labels)    # [bs, windows * 2 * K, embed_size]
        
        # 中心词与背景词应该同时出现，因此pos_dot的sigmoid结果应该趋于1
        pos_dot = torch.bmm(pos_embedding, input_embedding) # [batch_size, (window * 2), 1]
        pos_dot = pos_dot.squeeze(2) # [batch_size, (window * 2)]
        
        # 中心词与噪声词（负采样）不应该同时出现，因此pos_dot的sigmoid结果应该趋于0，
        # 但由于sigmoid函数的输出越接近1， logsigmoid的输出越接近0 
        # 因此多一个负号， 使得igmoid结果应该趋于1
        neg_dot = torch.bmm(neg_embedding, -input_embedding) # [batch_size, (window * 2 * K), 1]
        neg_dot = neg_dot.squeeze(2) # batch_size, (window * 2 * K)]
        
        # sigmoid函数的输出在为0-1之间， 则logsigmoid的输出全都小于0
        # 当sigmoid函数的输出越接近1， 则logsigmoid的输出越接近0
        # 可以理解为输出为1的损失小，为0的损失大
        log_pos = F.logsigmoid(pos_dot).sum(1) 
        log_neg = F.logsigmoid(neg_dot).sum(1)
        
        loss = log_pos + log_neg
        
        # logsigmoid的输出全都小于0， 如果要最小化loss，需要取负号
        return -loss

模型训练

data_path = 'text8/text8.train.txt'
batch_size = 32
lr = 0.2
epochs = 2
MAX_VOCAB_SIZE = 10000
EMBEDDING_SIZE = 100

my_loader = My_data_loader(data_path, batch_size)
model = Embedding_Model(MAX_VOCAB_SIZE, EMBEDDING_SIZE)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for e in range(epochs):
    for i, (input_labels, pos_labels, neg_labels) in enumerate(my_loader.loader):
        input_labels = input_labels.long()
        pos_labels = pos_labels.long()
        neg_labels = neg_labels.long()

        optimizer.zero_grad()
        loss = model(input_labels, pos_labels, neg_labels).mean()
        loss.backward()

        optimizer.step()

        if i % 100 == 0:
            print('epoch', e, 'iteration', i, loss.item())

测试词向量

写个函数，找出与某个词相近的一些词，比方说输入 good，他能帮我找出 nice，better，best 之类的

word2idx = my_loader.process_data.word2id
idx2word= my_loader.process_data.idx2word
embedding_weights = model.in_embed
def find_nearest(word):
    index = word2idx[word]
    embedding = embedding_weights(index)
    cos_dis = np.array([scipy.spatial.distance.cosine(e, embedding) for e in embedding_weights])
    return [idx2word[i] for i in cos_dis.argsort()[:10]]

for word in ["two", "america", "computer"]:
    print(word, find_nearest(word))
# 输出
two ['two', 'zero', 'four', 'one', 'six', 'five', 'three', 'nine', 'eight', 'seven']
america ['america', 'states', 'japan', 'china', 'usa', 'west', 'africa', 'italy', 'united', 'kingdom']
computer ['computer', 'machine', 'earth', 'pc', 'game', 'writing', 'board', 'result', 'code', 'website']