사실 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
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]
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)
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)
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
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]))
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')