Word2Vec

사실 Word2vec는 noise contrastive estimator (이하 NCE) loss를 사용한다. 아직 pytorch에서는 이 부분이 구현되어 있지 않고, 간단한 vocabulary이라서 그냥 softmax를 사용해서 이 부분을 구현하였다.

embedding이 2개이면, 단어에 따른 간단한 Classifiaction 문제로 볼 수 있기 때문에, 큰 무리는 없을 것이다.

※ 단, vocabulary수가 많아지면 학습속도를 높이기 위해서 NCE를 사용해야 할 것이다.


In [1]:
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data_utils

1. Dataset 준비


In [2]:
import numpy as np

word_pair = [['고양이', '흰'],
             ['고양이', '동물'],
             ['국화', '흰'],
             ['국화', '식물'],
             ['선인장', '초록'],
             ['선인장', '식물'],
             ['강아지', '검은'],
             ['강아지', '동물'],
             ['타조', '회색'],
             ['타조', '동물'],
             ['코끼리', '회색'],
             ['코끼리', '동물'],
             ['장미', '빨간'],
             ['장미', '식물'],
             ['자동차', '빨간'],
             ['그릇', '빨간'],
             ['민들레', '식물'],
             ['민들레', '흰']]

word_list = set(np.array(word_pair).flatten())
word_dict = {w: i for i, w in enumerate(word_list)}
skip_grams = [[word_dict[word[0]], word_dict[word[1]]] for word in word_pair]

Dataset Loader 설정


In [3]:
label = torch.LongTensor(skip_grams)[:, 0].contiguous()
context = torch.LongTensor(skip_grams)[:, 1].contiguous()
skip_grams_dataset = data_utils.TensorDataset(label, context)
train_loader = torch.utils.data.DataLoader(skip_grams_dataset, batch_size=8, shuffle=True)
test_loader = torch.utils.data.DataLoader(skip_grams_dataset, batch_size=1, shuffle=False)

2. 사전 설정

* model
* loss
* opimizer

In [4]:
class _model(nn.Module) :
    def __init__(self):
        super(_model, self).__init__()
        self.embedding = nn.Embedding(len(word_list), 2)
        self.linear = nn.Linear(2, len(word_list), bias=True)
        
    def forward(self, x):
        x = self.embedding(x)
        x = self.linear(x)
        return F.log_softmax(x)
    
model = _model()
loss_fn = nn.NLLLoss() 
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

3. Trainning loop

* (입력 생성)
* model 생성
* loss 생성
* zeroGrad
* backpropagation
* optimizer step (update model parameter)

In [5]:
model.train()
for epoch in range(100):
    for data, target in train_loader:
        data, target = Variable(data), Variable(target) #(입력 생성)
        output = model(data) # model 생성
        loss = F.nll_loss(output, target) #loss 생성
        optimizer.zero_grad() # zeroGrad
        loss.backward() # calc backward gradients
        optimizer.step() # update parameters

4. Predict & Evaluate


In [6]:
model.eval()

invDic = { i : w for w, i in word_dict.items()}
print('Input : true : pred')

for x, y in test_loader :
    x, y = Variable(x.squeeze()), y.squeeze()
    y_pred = model(x).max(1)[1].data[0][0]
    print('{:s} : {:s} : {:s}'.format(invDic[x.data[0]], invDic[y[0]], invDic[y_pred]))


Input : true : pred
고양이 : 흰 : 동물
고양이 : 동물 : 동물
국화 : 흰 : 식물
국화 : 식물 : 식물
선인장 : 초록 : 식물
선인장 : 식물 : 식물
강아지 : 검은 : 검은
강아지 : 동물 : 검은
타조 : 회색 : 회색
타조 : 동물 : 회색
코끼리 : 회색 : 회색
코끼리 : 동물 : 회색
장미 : 빨간 : 식물
장미 : 식물 : 식물
자동차 : 빨간 : 빨간
그릇 : 빨간 : 빨간
민들레 : 식물 : 식물
민들레 : 흰 : 식물

5. plot embedding space


In [7]:
import matplotlib.pyplot as plt
import matplotlib
%matplotlib inline 
matplotlib.rc('font', family="NanumGothic")

In [8]:
for i in label :
    x = Variable(torch.LongTensor([i]))
    fx, fy = model.embedding(x).squeeze().data
    plt.scatter(fx, fy)
    plt.annotate(invDic[i], xy=(fx, fy), xytext=(5, 2),
                 textcoords='offset points', ha='right', va='bottom')