本文主要是使用PyTorch復現word2vec論文
PyTorch中的nn.Embedding
實現關鍵是nn.Embedding()
這個API,首先看一下它的參數說明
[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-uWyBPAc4-1587546622043)(https://s1.ax1x.com/2020/04/12/GOaDhT.png#shadow)]
其中兩個必選參數num_embeddings
表示單詞的總數目,embedding_dim
表示每個單詞需要用什麼維度的向量表示。而nn.Embedding
權重的維度也是(num_embeddings, embedding_dim)
,默認是隨機初始化的
import torch
import torch.nn as nn
embeds = nn.Embedding(2, 5)
embeds.weight
# 輸出:
Parameter containing:
tensor([[-1.1454, 0.3675, -0.3718, 0.3733, 0.5979],
[-0.7952, -0.9794, 0.6292, -0.3633, -0.2037]], requires_grad=True)
如果使用與訓練好的詞向量,則採用
pretrained_weight = np.array(pretrained_weight)
embeds.weight.data.copy_(torch.from_numpy(pretrained_weight))
想要查看某個詞的詞向量,需要傳入這個詞在詞典中的index,並且這個index得是LongTensor型的
embeds = nn.Embedding(100, 10)
embeds(torch.LongTensor([50]))
# 輸出
tensor([[-1.9562e-03, 1.8971e+00, 7.0230e-01, -6.3762e-01, -1.9426e-01,
3.4200e-01, -2.0908e+00, -3.0827e-01, 9.6250e-01, -7.2700e-01]],
grad_fn=<EmbeddingBackward>)
過程詳解
具體的word2vec理論可以在我的這篇博客看到,這裏就不多贅述
下面說一下實現部分的細節
首先Embedding層輸入的shape是(batchsize, seq_len)
,輸出的shape是(batchsize, embedding_dim)
[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-BTe13Zws-1587546622046)(https://s1.ax1x.com/2020/04/12/GOImi8.png#shadow)]
上圖的流程是把文章中的單詞使用詞向量來表示
- 提取文章所有的單詞,把所有的單詞按照頻次降序排序(取前4999個,表示常出現的單詞。其餘所有單詞均用’'表示。所以一共有5000個單詞)
- 500個單詞使用one-hot編碼
- 通過訓練會生成一個的矩陣,每一行向量表示一個詞的詞向量。這裏的300是人爲指定,想要每個詞最終編碼爲詞向量的維度,你也可以設置成別的
這個矩陣如何獲得呢?在Skip-gram模型中,首先會隨機初始化這個矩陣,然後通過一層神經網絡來訓練。最終這個一層神經網絡的所有權重,就是要求的詞向量的矩陣
從上面的圖中看到,我們所學習的embedding層是一個訓練任務的一小部分,根據任務目標反向傳播,學習到embedding層裏的權重weight。這個weight是類似一種字典的存在,他能根據你輸入的one-hot向量查到相應的Embedding vector
Pytorch實現
導包
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as tud
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
random.seed(1)
np.random.seed(1)
torch.manual_seed(1)
C = 3 # context window
K = 15 # number of negative samples
epochs = 2
MAX_VOCAB_SIZE = 10000
EMBEDDING_SIZE = 100
batch_size = 32
lr = 0.2
上面的代碼我想應該沒有不明白的,C
就是論文中選取左右多少個單詞作爲背景詞。這裏我使用的是負採樣來近似訓練,K=15
表示隨機選取15個噪聲詞。MAX_VOCAB_SIZE=10000
表示這次實驗我準備訓練10000個詞的詞向量,但實際上我只會選出語料庫中出現次數最多的9999個詞,還有一個詞是<UNK>
用來表示所有的其它詞。每個詞的詞向量維度爲EMBEDDING_SIZE
語料庫下載地址:https://pan.baidu.com/s/10Bd3JxCCFTjBPNt0YROvZA 提取碼:81fo
[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-QUVrAW75-1587546622053)(https://s1.ax1x.com/2020/04/13/Gv38YD.png#shadow)]
文件中的內容是英文文本,去除了標點符號,每個單詞之間用空格隔開
讀取文本數據並處理
with open('text8.train.txt') as f:
text = f.read() # 得到文本內容
text = text.lower().split() # 分割成單詞列表
vocab_dict = dict(Counter(text).most_common(MAX_VOCAB_SIZE - 1)) # 得到單詞字典表,key是單詞,value是次數
vocab_dict['<UNK>'] = len(text) - np.sum(list(vocab_dict.values())) # 把不常用的單詞都編碼爲"<UNK>"
idx2word = [word for word in vocab_dict.keys()]
word2idx = {word:i for i, word in enumerate(idx2word)}
word_counts = np.array([count for count in vocab_dict.values()], dtype=np.float32)
word_freqs = word_counts / np.sum(word_counts)
word_freqs = word_freqs ** (3./4.)
最後一行代碼,word_freqs
存儲了每個單詞的頻率,然後又將所有的頻率變爲原來的0.75次方,這是因爲word2vec論文裏面推薦這麼做,當然你不改變這個值也沒什麼問題
[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-k2LpF4ig-1587546622056)(https://s1.ax1x.com/2020/04/13/GvGoQJ.png#shadow)]
實現DataLoader
接下來我們需要實現一個DataLoader,DataLoader可以幫助我們輕鬆打亂數據集,迭代的拿到一個mini-batch的數據等。一個DataLoader需要以下內容:
-
把所有word編碼成數字
-
保存vocabulary,單詞count、normalized word frequency
-
每個iteration sample一箇中心詞
-
根據當前的中心詞返回context單詞
-
根據中心詞sample一些negative單詞
-
返回單詞的counts
爲了使用DataLoader,我們需要定義以下兩個function
__len__()
:返回整個數據集有多少item__getitem__(idx)
:根據給定的idx返回一個item
這裏有一個好的tutorial介紹如何使用PyTorch DataLoader
class WordEmbeddingDataset(tud.Dataset):
def __init__(self, text, word2idx, idx2word, word_freqs, word_counts):
''' text: a list of words, all text from the training dataset
word2idx: the dictionary from word to index
idx2word: index to word mapping
word_freqs: the frequency of each word
word_counts: the word counts
'''
super(WordEmbeddingDataset, self).__init__() # #通過父類初始化模型,然後重寫兩個方法
self.text_encoded = [word2idx.get(word, word2idx['<UNK>']) for word in text] # 把單詞數字化表示。如果不在詞典中,也表示爲unk
self.text_encoded = torch.LongTensor(self.text_encoded) # nn.Embedding需要傳入LongTensor類型
self.word2idx = word2idx
self.idx2word = idx2word
self.word_freqs = torch.Tensor(word_freqs)
self.word_counts = torch.Tensor(word_counts)
def __len__(self):
return len(self.text_encoded) # 返回所有單詞的總數,即item的總數
def __getitem__(self, idx):
''' 這個function返回以下數據用於訓練
- 中心詞
- 這個單詞附近的positive word
- 隨機採樣的K個單詞作爲negative word
'''
center_words = self.text_encoded[idx] # 取得中心詞
pos_indices = list(range(idx - C, idx)) + list(range(idx + 1, idx + C + 1)) # 先取得中心左右各C個詞的索引
pos_indices = [i % len(self.text_encoded) for i in pos_indices] # 爲了避免索引越界,所以進行取餘處理
pos_words = self.text_encoded[pos_indices] # tensor(list)
neg_words = torch.multinomial(self.word_freqs, K * pos_words.shape[0], True)
# torch.multinomial作用是對self.word_freqs做K * pos_words.shape[0]次取值,輸出的是self.word_freqs對應的下標
# 取樣方式採用有放回的採樣,並且self.word_freqs數值越大,取樣概率越大
# 每採樣一個正確的單詞(positive word),就採樣K個錯誤的單詞(negative word),pos_words.shape[0]是正確單詞數量
return center_words, pos_words, neg_words
每一行代碼詳細的註釋都寫在上面了,其中有一行代碼需要特別說明一下,就是註釋了tensor(list)的那一行,因爲text_encoded
本身是個tensor,而傳入的pos_indices
是一個list。下面舉個例子就很好理解這句代碼的作用了
a = torch.tensor([2, 3, 3, 8, 4, 6, 7, 8, 1, 3, 5, 0], dtype=torch.long)
b = [2, 3, 5, 6]
print(a[b])
# tensor([3, 8, 6, 7])
通過下面兩行代碼即可得到DataLoader
dataset = WordEmbeddingDataset(text, word2idx, idx2word, word_freqs, word_counts)
dataloader = tud.DataLoader(dataset, batch_size, shuffle=True)
可以隨便打印一下看看
next(iter(dataset))
'''
(tensor(4813),
tensor([ 50, 9999, 393, 3139, 11, 5]),
tensor([ 82, 0, 2835, 23, 328, 20, 2580, 6768, 34, 1493, 90, 5,
110, 464, 5760, 5368, 3899, 5249, 776, 883, 8522, 4093, 1, 4159,
5272, 2860, 9999, 6, 4880, 8803, 2778, 7997, 6381, 264, 2560, 32,
7681, 6713, 818, 1219, 1750, 8437, 1611, 12, 42, 24, 22, 448,
9999, 75, 2424, 9970, 1365, 5320, 878, 40, 2585, 790, 19, 2607,
1, 18, 3847, 2135, 174, 3446, 191, 3648, 9717, 3346, 4974, 53,
915, 80, 78, 6408, 4737, 4147, 1925, 4718, 737, 1628, 6160, 894,
9373, 32, 572, 3064, 6, 943]))
'''
定義PyTorch模型
class EmbeddingModel(nn.Module):
def __init__(self, vocab_size, embed_size):
super(EmbeddingModel, self).__init__()
self.vocab_size = vocab_size
self.embed_size = embed_size
self.in_embed = nn.Embedding(self.vocab_size, self.embed_size)
self.out_embed = nn.Embedding(self.vocab_size, self.embed_size)
def forward(self, input_labels, pos_labels, neg_labels):
''' input_labels: center words, [batch_size]
pos_labels: positive words, [batch_size, (window_size * 2)]
neg_labels:negative words, [batch_size, (window_size * 2 * K)]
return: loss, [batch_size]
'''
input_embedding = self.in_embed(input_labels) # [batch_size, embed_size]
pos_embedding = self.out_embed(pos_labels)# [batch_size, (window * 2), embed_size]
neg_embedding = self.out_embed(neg_labels) # [batch_size, (window * 2 * K), embed_size]
input_embedding = input_embedding.unsqueeze(2) # [batch_size, embed_size, 1]
pos_dot = torch.bmm(pos_embedding, input_embedding) # [batch_size, (window * 2), 1]
pos_dot = pos_dot.squeeze(2) # [batch_size, (window * 2)]
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)]
log_pos = F.logsigmoid(pos_dot).sum(1) # .sum()結果只爲一個數,.sum(1)結果是一維的張量
log_neg = F.logsigmoid(neg_dot).sum(1)
loss = log_pos + log_neg
return -loss
def input_embedding(self):
return self.in_embed.weight.numpy()
這裏爲什麼要分兩個embedding層來訓練?很明顯,對於任一一個詞,它既有可能作爲中心詞出現,也有可能作爲背景詞出現,所以每個詞需要用兩個向量去表示。in_embed
訓練出來的權重就是每個詞作爲中心詞的權重。out_embed
訓練出來的權重就是每個詞作爲背景詞的權重。那麼最後到底用什麼向量來表示一個詞呢?是中心詞向量?還是背景詞向量?按照Word2Vec論文所寫,推薦使用中心詞向量,所以這裏我最後返回的是in_embed.weight
。如果上面我說的你不太明白,可以看我之前的Word2Vec詳解
bmm(a, b)
,batch matrix multiply。函數中的兩個參數a,b都是維度爲3的tensor,並且這兩個tensor的第一個維度必須相同,後面兩個維度必須滿足矩陣乘法的要求
batch1 = torch.randn(10, 3, 4)
batch2 = torch.randn(10, 4, 5)
res = torch.bmm(batch1, batch2)
print(res.size())
# torch.Size([10, 3, 5])
訓練模型
for e in range(1):
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()
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())
embedding_weights = model.input_embeddings()
torch.save(model.state_dict(), "embedding-{}.th".format(EMBEDDING_SIZE))
如果沒有GPU,訓練時間可能比較長
詞向量應用
我們可以寫個函數,找出與某個詞相近的一些詞,比方說輸入good,他能幫我找出nice,better,best之類的
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 [idx_to_word[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']
nn.Linear VS. nn.Embedding
Word2Vec論文中給出的架構其實就一個單層神經網絡,那麼爲什麼直接用nn.Linear()
來訓練呢?nn.Linear()
不是也能訓練出一個weight嗎?
答案是可以的,當然可以直接使用nn.Linear()
,只不過輸入要改爲one-hot Encoding,而不能像nn.Embedding()
這種方式直接傳入一個index。還有就是需要設置bias=False
,因爲我們只需要訓練一個權重矩陣,不訓練偏置
這裏給出一個使用單層神經網絡來訓練Word2Vec的博客