一個基於PyTorch實現的Glove詞向量的實例

簡介

詞向量技術，也稱爲詞嵌入技術(word-embedding)，是一種將高維稀疏的向量壓縮到低維稠密向量的技術。常見於自然語言處理領域對單詞的預處理過程，例如將單詞的one-hot向量是高維稀疏的，不但佔用大量空間，而且向量之間提供的信息很少。但經過詞嵌入技術生成的向量，不但是低維的並且包含了很多詞語間的語義、語法的信息。因此在NLP中，往往會先針對語料生成相應的詞向量，然後再把詞向量喂入具體任務的神經網絡。

Glove詞向量是常用的詞向量之一，由斯坦福大學發佈，原版代碼是C實現的，有興趣的讀者可以前往Glove官網下載。本文是根據Glove論文的原理，利用深度學習庫PyTorch實現的一個簡易版本，僅僅用作學習的示例。

相關的庫

1、python3.6
2、torch
3、numpy
4、scipy
5、text8英文語料

Glove原理

在開始寫代碼之前，我們對Glove原理進行一次快速介紹。Glove詞向量核心思想是利用詞與詞之間的共現次數來進行訓練的，所謂共現就是在一個上下文窗口大小範圍內兩個詞語同時出現的次數，把上下文窗口從語料庫的頭到尾遍歷一次，就得到了一個全局的共現矩陣，然後再基於這個共現矩陣進行詞向量的訓練。

下面，我們來看看需要優化的目標函數：
$J = \displaystyle\sum_{i,j}^N f(X_{i,j})(v_i^Tw_j+b_i+b_j-log(X_{i,j}))^2$
其中，$N$表示詞彙表大小，$X_{i,j}$表示詞語i和詞語j的共現次數，$v_i$表示中心詞i對應的詞向量，$w_j$表示上下文詞j對應的上下文向量。$bi、bj$是偏置項，而$f(·)$表示權重函數，該函數會抑制共現次數過高的詞對造成的影響。從目標函數可以看出，模型會訓練出兩個向量$v、w$，原文作者建議將這二者的和作爲最後的詞向量表示。

實現

1、加載語料庫進行預處理
這一步主要是把語料加載到內存中，生成詞彙表等。我們直接來看代碼：

def getCorpus(filetype, size):
    if filetype == 'dev':
        filepath = '../corpus/text8.dev.txt'
    elif filetype == 'test':
        filepath = '../corpus/text8.test.txt'
    else:
        filepath = '../corpus/text8.train.txt'

    with open(filepath, "r") as f:
        text = f.read()
        text = text.lower().split()
        text = text[: min(len(text), size)]
        vocab_dict = dict(Counter(text).most_common(MAX_VOCAB_SIZE - 1))
        vocab_dict['<unk>'] = len(text) - sum(list(vocab_dict.values()))
        idx_to_word = list(vocab_dict.keys())
        word_to_idx = {word:ind for ind, word in enumerate(idx_to_word)}
        word_counts = np.array(list(vocab_dict.values()), dtype=np.float32)
        word_freqs = word_counts / sum(word_counts)
        print("Words list length:{}".format(len(text)))
        print("Vocab size:{}".format(len(idx_to_word)))
    return text, idx_to_word, word_to_idx, word_counts, word_freqs

內容很簡單，因爲語料庫也不大，所以能直接加載到內存中。根據語料庫內的單詞，挑選出前2000個頻率最高的詞語構成詞彙表，剩下的詞語則用UNK代替。然後生成一個單詞-序號、序號-單詞的映射表、詞頻列表等，以便後續的查詢操作。

2、計算詞共現矩陣
詞共現矩陣的計算方式如下：給出一個句子如"china is a wonderful country and everyone is kind"。假設中心詞是wonderful，窗口大小是2，那麼我我們增加"wonderful,is"、“wonderful,a”、“wonderful,country”、"wonderful,and"的共現次數，即：$X_{wonderful,is}$+=1，$X_{wonderful,a}$+=1、$X_{wonderful,country}$+=1等。接着把窗口向右移動，統計下一個中心詞-上下文詞的共現次數，直到遍歷完整個語料庫。

def buildCooccuranceMatrix(text, word_to_idx):
    vocab_size = len(word_to_idx)
    maxlength = len(text)
    text_ids = [word_to_idx.get(word, word_to_idx["<unk>"]) for word in text]
    cooccurance_matrix = np.zeros((vocab_size, vocab_size), dtype=np.int32)
    print("Co-Matrix consumed mem:%.2fMB" % (sys.getsizeof(cooccurance_matrix)/(1024*1024)))
    for i, center_word_id in enumerate(text_ids):
        window_indices = list(range(i - WINDOW_SIZE, i)) + list(range(i + 1, i + WINDOW_SIZE + 1))
        window_indices = [i % maxlength for i in window_indices]
        window_word_ids = [text_ids[index] for index in window_indices]
        for context_word_id in window_word_ids:
            cooccurance_matrix[center_word_id][context_word_id] += 1
        if (i+1) % 1000000 == 0:
            print(">>>>> Process %dth word" % (i+1))
    print(">>>>> Save co-occurance matrix completed.")
    return cooccurance_matrix

3、計算權重函數矩陣
由於經常出現的詞語實際上不會提供太多的信息，比如to、the之類的停用詞，此時他們與別的詞的共現頻率就會很高，對訓練詞向量是一個阻礙，因此需要抑制共現頻率過高造成的影響，使其在一個合理的範圍內。根據原論文所述，其權重函數形式如下：
$f(x) = \begin{cases} (x/xmax)^{0.75} &\text{if } x < xmax \\ 1 &\text{if } x ≥ xmax \end{cases}$
其中，xmax=100.0

對於$X_{i,j}$可以分別計算出一個對應的$f(X_{i,j})$，由此可以形成一個權重矩陣。實際上權重的計算可以放到前向傳播的時候實時計算出來，不需要這樣一個矩陣，但筆者這裏爲了節省前向傳播的時間，提前計算出來並保存到矩陣內，等需要用的時候直接查找即可。代碼如下所示：

def buildWeightMatrix(co_matrix):
    xmax = 100.0
    weight_matrix = np.zeros_like(co_matrix, dtype=np.float32)
    print("Weight-Matrix consumed mem:%.2fMB" % (sys.getsizeof(weight_matrix) / (1024 * 1024)))
    for i in range(co_matrix.shape[0]):
        for j in range(co_matrix.shape[1]):
            weight_matrix[i][j] = math.pow(co_matrix[i][j] / xmax, 0.75) if co_matrix[i][j] < xmax else 1
        if (i+1) % 1000 == 0:
            print(">>>>> Process %dth weight" % (i+1))
    print(">>>>> Save weight matrix completed.")
    return weight_matrix

4、創建DataLoader
DataLoader是PyTorch的一個數據加載器，利用它可以很方便地把訓練集分批、打亂訓練集等，我們新建一個類WordEmbeddingDataset繼承自torch.untils.data.Dataset，代碼如下圖所示：

class WordEmbeddingDataset(tud.Dataset):
    def __init__(self, co_matrix, weight_matrix):
        self.co_matrix = co_matrix
        self.weight_matrix = weight_matrix
        self.train_set = []

        for i in range(self.weight_matrix.shape[0]):
            for j in range(self.weight_matrix.shape[1]):
                if weight_matrix[i][j] != 0:
                    # 這裏對權重進行了篩選，去掉權重爲0的項 
                    # 因爲共現次數爲0會導致log(X)變成nan
                    self.train_set.append((i, j))   

    def __len__(self):
        '''
        必須重寫的方法
        :return: 返回訓練集的大小
        '''
        return len(self.train_set)

    def __getitem__(self, index):
        '''
        必須重寫的方法
        :param index:樣本索引 
        :return: 返回一個樣本
        '''
        (i, j) = self.train_set[index]
        return i, j, torch.tensor(self.co_matrix[i][j], dtype=torch.float), self.weight_matrix[i][j]

可以看出，返回的樣本包括了詞語i和j在詞彙表中的需要，以及他們的共現頻率和權重，這樣就可以直接用於前向傳播的計算中。

5、創建訓練模型
下面來創建我們的訓練模型，在pytorch中創建模型很簡單，只需要繼承nn.Module，然後重寫前向傳播函數即可，由於pytorch有自動求導機制，所以我們不需要擔心反向傳播的問題。

class GloveModelForBGD(nn.Module):
    def __init__(self, vocab_size, embed_size):
        super().__init__()
        self.vocab_size = vocab_size
        self.embed_size = embed_size
        
        #聲明v和w爲Embedding向量
        self.v = nn.Embedding(vocab_size, embed_size)
        self.w = nn.Embedding(vocab_size, embed_size)
        self.biasv = nn.Embedding(vocab_size, 1)
        self.biasw = nn.Embedding(vocab_size, 1)
        
        #隨機初始化參數
        initrange = 0.5 / self.embed_size
        self.v.weight.data.uniform_(-initrange, initrange)
        self.w.weight.data.uniform_(-initrange, initrange)

    def forward(self, i, j, co_occur, weight):
    	#根據目標函數計算Loss值
        vi = self.v(i)	#分別根據索引i和j取出對應的詞向量和偏差值
        wj = self.w(j)
        bi = self.biasv(i)
        bj = self.biasw(j)

        similarity = torch.mul(vi, wj)
        similarity = torch.sum(similarity, dim=1)

        loss = similarity + bi + bj - torch.log(co_occur)
        loss = 0.5 * weight * loss * loss

        return loss.sum().mean()

    def gloveMatrix(self):
        '''
        獲得詞向量，這裏把兩個向量相加作爲最後的詞向量
        :return: 
        '''
        return self.v.weight.data.numpy() + self.w.weight.data.numpy()

6、訓練
下面是對上文的幾個點進行串聯，最後進行訓練，具體的可以參考代碼：

EMBEDDING_SIZE = 50		#50個特徵
MAX_VOCAB_SIZE = 2000	#詞彙表大小爲2000個詞語
WINDOW_SIZE = 5			#窗口大小爲5

NUM_EPOCHS = 10			#迭代10次
BATCH_SIZE = 10			#一批有10個樣本
LEARNING_RATE = 0.05	#初始學習率
TEXT_SIZE = 20000000	#控制從語料庫讀取語料的規模

text, idx_to_word, word_to_idx, word_counts, word_freqs = getCorpus('train', size=TEXT_SIZE)    #加載語料及預處理
co_matrix = buildCooccuranceMatrix(text, word_to_idx)    #構建共現矩陣
weight_matrix = buildWeightMatrix(co_matrix)             #構建權重矩陣
dataset = WordEmbeddingDataset(co_matrix, weight_matrix) #創建dataset
dataloader = tud.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
model = GloveModelForBGD(MAX_VOCAB_SIZE, EMBEDDING_SIZE) #創建模型
optimizer = torch.optim.Adagrad(model.parameters(), lr=LEARNING_RATE) #選擇Adagrad優化器

print_every = 10000
save_every = 50000
epochs = NUM_EPOCHS
iters_per_epoch = int(dataset.__len__() / BATCH_SIZE)
total_iterations = iters_per_epoch * epochs
print("Iterations: %d per one epoch, Total iterations: %d " % (iters_per_epoch, total_iterations))
start = time.time()
for epoch in range(epochs):
	loss_print_avg = 0
	iteration = iters_per_epoch * epoch
	for i, j, co_occur, weight in dataloader:
		iteration += 1
        optimizer.zero_grad()   #每一批樣本訓練前重置緩存的梯度
        loss = model(i, j, co_occur, weight)    #前向傳播
        loss.backward()     #反向傳播
        optimizer.step()    #更新梯度
        loss_print_avg += loss.item()
torch.save(model.state_dict(), WEIGHT_FILE)

實驗結果

1、利用餘弦相似度衡量向量之間的相似性
按照上述代碼進行訓練後，我們執行以下代碼來看看訓練得到的模型的效果怎麼樣。我們利用餘弦相似度來衡量兩個向量是否相似，如果兩個詞語是同屬一類的詞語，那麼它們的向量的距離應該很近的。我們定義以下的函數：

def find_nearest(word, embedding_weights):
    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]]#找到前10個最相近詞語

然後列舉幾個詞語，看看跟他們最相似的幾個詞語是什麼，運行以下代碼：

glove_matrix = model.gloveMatrix()
for word in ["good", "one", "green", "like", "america", "queen", "better", "paris", "work", "computer", "language"]:
	print(word, find_nearest(word, glove_matrix))

運行結果如下：

可以看出，在某些詞語比如數字、顏色、計算機、形容詞比較級等，取得了不錯的效果。

2、利用向量之間的運算考察詞向量的關聯性
詞向量有另外一個特性，即“man to king is woman to ?”，向量化表述就是$v(candidate) = v(king)-v(man)+v(woman)$，也即是候選詞的詞向量可以由另外三個詞向量通過加減運算得到。那麼，顯然上述的候選詞應該是“queen”。我們利用這個特性來舉幾個例子：

def findRelationshipVector(word1, word2, word3):
    word1_idx = word_to_idx[word1]
    word2_idx = word_to_idx[word2]
    word3_idx = word_to_idx[word3]
    embedding = glove_matrix[word2_idx] - glove_matrix[word1_idx] + glove_matrix[word3_idx]
    cos_dis = np.array([scipy.spatial.distance.cosine(e, embedding) for e in glove_matrix])
    for i in cos_dis.argsort()[:5]:
        print("{} to {} as {} to {}".format(word1, word2, word3, idx_to_word[i]))

findRelationshipVector('man', 'king', 'woman')
findRelationshipVector('america', 'washington', 'france')
findRelationshipVector('good', 'better', 'bad')

運行以上代碼，可以得到以下結果：

可以看出，結果一般般，但還是能看得到模型在努力地篩選出了更爲相關的信息，只不過準確度需要再進一步提高。

3、利用SVD進行詞向量的降維並可視化顯示
SVD技術常用於將高維向量降維，我們將詞向量降到二維，然後把降維後的向量畫到畫板上，看看這些詞向量有沒有很好地聚合在一起。如果聚合在一起了，表示詞向量是相近的，也即是模型有比較好地學習到了詞語之間的語義、語法關係。運行以下代碼並觀察結果：

#數據降維以及可視化
candidate_words = ['one','two','three','four','five','six','seven','eight','night','ten','color','green','blue','red','black',                      'man','woman','king','queen','wife','son','daughter','brown','zero','computer','hardware','software','system','program',
   'america','china','france','washington','good','better','bad']
candidate_indexes = [word_to_idx[word] for word in candidate_words]
choosen_indexes = candidate_indexes
choosen_vectors = [glove_matrix[index] for index in choosen_indexes]

U, S, VH = np.linalg.svd(choosen_vectors, full_matrices=False)
for i in range(len(choosen_indexes)):
plt.text(U[i, 0], U[i, 1], idx_to_word[choosen_indexes[i]])

coordinate = U[:, 0:2]
plt.xlim((np.min(coordinate[:, 0]) - 0.1, np.max(coordinate[:, 0]) + 0.1))
plt.ylim((np.min(coordinate[:, 1]) - 0.1, np.max(coordinate[:, 1]) + 0.1))
plt.show()

從結果圖可以觀察到，同類型語義相近的詞語基本都靠得很近。

結論

通過對實驗結果的觀察，可以發現訓練得到的Glove模型效果還不錯，由於筆者設備的限制，僅僅對1500萬個詞語的語料庫以及2000個詞語的詞彙表，採取了50個特徵值，進行了10個輪次的訓練。如果語料庫更大、詞彙表更大以及特徵數量更多，相信訓練出來的詞向量將會有更好的效果。由於本文僅僅是作爲一個學習的樣例，所以Glove的訓練到此爲止，謝謝各位的閱讀~

項目代碼

最後貼出筆者的項目代碼地址，包括了語料庫以及模型等，有需要的同學可以自行下載：PyTorchGlove