class Word2Vec:
def __init__(self, input_file_name, output_file_name):
self.output_file_name = output_file_name
self.data = InputData(input_file_name, MIN_COUNT)
self.model = SkipGramModel(self.data.word_count, EMB_DIMENSION)
self.lr = LR
self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr)
def train(self):
print("SkipGram Training......")
pairs_count = self.data.evaluate_pairs_count(WINDOW_SIZE)
print("pairs_count", pairs_count)
batch_count = pairs_count / BATCH_SIZE
print("batch_count", batch_count)
process_bar = tqdm(range(int(batch_count)))
for i in process_bar:
pos_pairs = self.data.get_batch_pairs(BATCH_SIZE, WINDOW_SIZE)
pos_pairs, neg_pairs = self.data.get_pairs(pos_pairs)
pos_u = [pair[0] for pair in pos_pairs]
pos_v = [int(pair[1]) for pair in pos_pairs]
neg_u = [pair[0] for pair in neg_pairs]
neg_v = [int(pair[1]) for pair in neg_pairs]
self.optimizer.zero_grad()
loss = self.model.forward(pos_u, pos_v, neg_u, neg_v)
loss.backward()
self.optimizer.step()
if i * BATCH_SIZE % 100000 == 0:
self.lr = self.lr * (1.0 - 1.0 * i / batch_count)
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.lr
self.model.save_embedding(self.data.id2word_dict,
self.output_file_name)
class Word2Vec:
def __init__(self, input_file_name, output_file_name):
self.output_file_name = output_file_name
self.data = InputData(input_file_name, MIN_COUNT)
self.model = SkipGramModel(self.data.word_count, EMB_DIMENSION).cuda()
self.lr = LR
self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr)
def train(self):
for _ in range(1, EPOCH + 1):
print("SkipGram Training......")
pairs_count = self.data.evaluate_pairs_count(WINDOW_SIZE)
print("pairs_count", pairs_count)
batch_count = pairs_count / BATCH_SIZE
print("batch_count", batch_count)
process_bar = tqdm(range(int(batch_count)))
for i in process_bar:
pos_pairs = self.data.get_batch_pairs(BATCH_SIZE, WINDOW_SIZE)
pos_pairs, neg_pairs = self.data.get_pairs(pos_pairs)
pos_u = [pair[0] for pair in pos_pairs]
pos_v = [int(pair[1]) for pair in pos_pairs]
neg_u = [pair[0] for pair in neg_pairs]
neg_v = [int(pair[1]) for pair in neg_pairs]
self.optimizer.zero_grad()
loss = self.model.forward(pos_u, pos_v, neg_u, neg_v)
loss.backward()
self.optimizer.step()
if i * BATCH_SIZE % 100000 == 0:
self.lr = self.lr * (1.0 - 1.0 * i / batch_count)
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.lr
process_bar.set_postfix(loss=loss.data.cpu().numpy())
process_bar.update()
print('\n')
torch.save(self.model.state_dict(), "../results/url_with_location_skipgram_hs_wyz.pkl")
self.model.save_embedding(self.data.id2word_dict, self.output_file_name)