教程:https://www.bilibili.com/video/BV1vz4y1R7Mm?p=2
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
USE_CUDA = torch.cuda.is_available()
random.seed(1)
np.random.seed(1)
torch.manual_seed(1)
if USE_CUDA:
torch.cuda.manual_seed(1)
C = 3 # content window
K = 100 # number of negtive samples
NUM_EPOCHS = 2
MAX_VOCAB_SIZE = 30000
BATCH_SIZE = 128
LEARNING_RATE = 0.2
EMBEDDING_SIZE = 100
def word_tokenize(text):
return text.split()
- 从文本文件中读取所有的文字,通过这些文本创建一个vocabulary
- 由于单词数量可能太大,我们只选取最常见的MAX_VOCAB_SIZE个单词
- 添加一个UNK单词表示所有不常见的单词
- 我们需要记录单词到index的mapping, 以及index到单词的mapping, 单词的count, 单词的normalized frequency, 单词总数
with open('text8/text8.train.txt', 'r') as fin:
text = fin.read()
text = word_tokenize(text)
vocab = dict(Counter(text).most_common(MAX_VOCAB_SIZE-1))
vocab['<unk>'] = len(text) - np.sum(list(vocab.values()))
idx_to_word = [word for word in vocab.keys()]
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_counts = np.array([count for count in vocab.values()], dtype=np.float32)
word_freqs = word_counts / np.sum(word_counts)
word_freqs = word_freqs ** (0.75)
word_freqs = word_freqs / np.sum(word_freqs)
实现DataLoader
返回batch的数据
class WordEmbeddingDataset(tud.Dataset):
def __init__(self, text, word_to_idx, idx_to_word, word_freqs, word_counts):
super()
self.text_encoded = [word_to_idx.get(word, word_to_idx['<unk>']) for word in text]
self.text_encoded = torch.LongTensor(self.text_encoded)
self.word_to_idx = word_to_idx
self.word_freqs = torch.Tensor(word_freqs)
self.word_counts = torch.Tensor(word_counts)
def __len__(self):
return len(self.text_encoded)
def __getitem__(self, idx):
center_word = self.text_encoded[idx]
pos_indices = list(range(idx-C, idx)) + list(range(idx+1, idx+C+1)) # window内的单词(写错了)
pos_indices = [i % len(self.text_encoded) for i in pos_indices] # 取余,放置超出长度
pos_words = self.text_encoded[pos_indices] # 周围单词
neg_words = torch.multinomial(self.word_freqs, K * pos_words.shape[0], True) # 负采样单词
# Returns a tensor where each row contains num_samples indices sampled from
# the multinomial probability distribution located in the corresponding row
# of tensor input. (多项式概率分布)
return center_word, pos_words, neg_words
dataset = WordEmbeddingDataset(text, word_to_idx, idx_to_word, word_freqs, word_counts)
dataloader = tud.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
dataset.word_freqs.shape
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(vocab_size, embed_size)
self.out_embed = nn.Embedding(vocab_size, embed_size)
def forward(self, input_labels, pos_labels, neg_labels):
# input_label: [batch_size]
# pos_labels: [batch_size, (window_size * 2)]
# neg_labels: [batch_size, (window_size * 2 * K)]
input_embedding = self.in_embed(input_labels) # [batch_size, embed_size]
pos_embedding = self.out_embed(pos_labels) # [batch_size, (window_size*2), embed_size]
neg_embedding = self.out_embed(neg_labels) # [batch_size, (window_size*2)*K, embed_size]
input_embedding = input_embedding.unsqueeze(2) # [batch_size, embed_size, 1]
pos_dot = torch.bmm(pos_embedding, input_embedding).squeeze(2)
# documention of torch.bmm: batch matrix-matrix product
# Performs a batch matrix-matrix product of matrices stored in input and mat2.
# if A.shape=[batch_size, n, m], B.shape=[batch_size, m, p],
# then torch.bmm(A, B).shape is [batch_size, n, p]
neg_dot = torch.bmm(neg_embedding, -input_embedding).squeeze(2)
log_pos = F.logsigmoid(pos_dot).sum(1)
log_neg = F.logsigmoid(neg_dot).sum(1)
loss= log_pos+log_neg
return -loss
def input_embeddings(self):
return self.in_embed.weight.data.cpu().numpy()
model = EmbeddingModel(MAX_VOCAB_SIZE, EMBEDDING_SIZE)
if USE_CUDA:
model = model.cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)
for e in range(NUM_EPOCHS):
for i, (input_labels, pos_labels, neg_labels) in enumerate(dataloader):
# print(input_labels, pos_labels, neg_labels)
# if i > 5: break
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:
print("epoch", e, "iteration", i, loss.item())