embeddings_pytorch.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Oct  4 13:15:27 2018

@author: dorgham
"""

import torch
import torch.nn as nn
from torch.nn import functional as F
import torch.optim as optim
import helpers
import numpy as np
import random as rn
import argparse
import sys



class Path2VecModel(nn.Module):
    
    def __init__(self, vocab_size, embedding_dim):
        super(Path2VecModel, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        

    def forward(self, inputs):
        embed1 = self.embeddings(inputs[0])
        embed2 = self.embeddings(inputs[1])
        #normalize the vectors before the dot product so that dot product is the cosine proximity between the two vectors
        embed1 = embed1 / embed1.norm(2, 2, True).clamp(min=1e-12).expand_as(embed1)
        embed2 = embed2 / embed2.norm(2, 2, True).clamp(min=1e-12).expand_as(embed2)
        out = torch.sum(embed1*embed2, dim=2)
        
        return out
    
    
    
def custom_loss(y_pred, y_true, reg_1_output, reg_2_output, use_neighbors, beta=0.01, gamma=0.01):
    if use_neighbors:
        alpha = 1 - (beta+gamma)
        m_loss = alpha * F.mse_loss(y_pred, y_true, reduction='elementwise_mean')
        
        m_loss -= beta * reg_1_output
        m_loss -= gamma * reg_2_output
    else:
        m_loss = F.mse_loss(y_pred, y_true, reduction='elementwise_mean')

    return m_loss


def load_training_data(trainfile, vocab_file):
    wordpairs = helpers.Wordpairs(trainfile)
    
    if not vocab_file:
        print('Building vocabulary from the training set...', file=sys.stderr)
        no_train_pairs, vocab_dict, inverted_vocabulary = helpers.build_vocabulary(wordpairs)
        print('Building vocabulary finished', file=sys.stderr)
    else:
        vocabulary_file = vocab_file  # JSON file with the ready-made vocabulary
        print('Loading vocabulary from file', vocabulary_file, file=sys.stderr)
        vocab_dict, inverted_vocabulary = helpers.vocab_from_file(vocabulary_file)
        print('Counting the number of pairs in the training set...')
        no_train_pairs = 0
        for line in wordpairs:
            no_train_pairs += 1
        print('Number of pairs in the training set:', no_train_pairs)
    
    return wordpairs, vocab_dict


def save_embeddings(filename, model, vocab_dict):
    # Saving the resulting vectors
    embeddings = model.state_dict()['embeddings.weight']
    if torch.cuda.is_available():
        embeddings = embeddings.cpu()
    helpers.save_word2vec_format(filename, vocab_dict, embeddings.numpy())
    
 
def run(trainfile, vocab_file, embedding_dimension, batch_size, learn_rate, neighbors_count, negative, run_name, 
        l1_factor, beta, gamma, fix_seeds, epochs, use_neighbors, regularize):    
    if fix_seeds:
        #fix seeds for repeatability of experiments
        np.random.seed(42)
        rn.seed(12345)
        torch.manual_seed(1)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(1)
    
    wordpairs, vocab_dict = load_training_data(trainfile, vocab_file)
    
    print('Retreiving neighbors of training samples...')
    helpers.build_connections(vocab_dict)
    
    vocab_size = len(vocab_dict)
    
    #instantiate the model
    model = Path2VecModel(vocab_size, embedding_dimension)
    #use GPU if available
    if torch.cuda.is_available():
        model.cuda()
        torch.cuda.manual_seed(1)
        print("Using GPU...")
    
    optimizer = optim.Adam(model.parameters(), lr=learn_rate)
    
    print('Model name and layers:')
    print(model)
    
    #begin the training..
    for epoch in range(epochs):
        print('Epoch #', epoch+1)
        total_loss, n_batches = 0, 0
        batchGenerator = helpers.batch_generator_2(wordpairs, vocab_dict, vocab_size, negative, batch_size)
        for batch in batchGenerator:
            n_batches +=1
            l1_reg_term = 0
            inputs, targets = batch
            target_tensor = torch.from_numpy(targets).float()
            
            input_var = torch.Tensor([inputs[0], inputs[1]]).long()
            if torch.cuda.is_available():
                input_var = input_var.cuda()
                target_tensor = target_tensor.cuda()
    
            
            model.zero_grad()
            #do the forward pass
            similarity_pred = model(input_var)

            if use_neighbors:
                #get only the positive samples because the batch variable contains the generated negatives as well
                positive_samples = helpers.get_current_positive_samples()
                inputs_list = [[], []]
                for word_idx in positive_samples[0]:
                    neighbors = helpers.get_node_neighbors(word_idx)
                    for neighbor in neighbors:
                        inputs_list[0].append([word_idx])
                        inputs_list[1].append([neighbor])
                        
                input_var = torch.Tensor(inputs_list).long()
                if torch.cuda.is_available():
                    input_var = input_var.cuda()
                    
                reg1_dot_prod = model(input_var)
                reg1_output = torch.sum(reg1_dot_prod) / len(reg1_dot_prod)

                inputs_list = [[], []]
                for word_idx in positive_samples[1]: #context words
                    neighbors = helpers.get_node_neighbors(word_idx)
                    for neighbor in neighbors:
                        inputs_list[0].append([word_idx])
                        inputs_list[1].append([neighbor])
                        
                input_var = torch.Tensor(inputs_list).long()
                if torch.cuda.is_available():
                    input_var = input_var.cuda()
                    
                reg2_dot_prod = model(input_var)
                reg2_output = torch.sum(reg2_dot_prod) / len(reg2_dot_prod)
                   
            # Compute the loss function. 
            loss = custom_loss(similarity_pred, target_tensor, reg1_output, reg2_output, use_neighbors, beta, gamma)
            if regularize == True:
                for param in model.parameters():
                    l1_reg_term += torch.norm(param, 1)
                loss += l1_factor * l1_reg_term
            
            # Do the backward pass and update the gradient
            loss.backward()
            optimizer.step()
    
            # normalize the loss per batch size
            total_loss += loss.item() / len(inputs[0])
            
        print('Total loss = ', total_loss/n_batches)
    
    
    train_name = trainfile.split('.')[0] + '_embeddings_vsize' + str(embedding_dimension) +'_bsize' + str(batch_size) \
                 + '_lr' + str(learn_rate).split('.')[-1]+'_nn-'+str(args.use_neighbors)+str(args.neighbor_count)+\
                 '_reg-'+str(args.regularize)
    filename = train_name + '_' + run_name + '.vec.gz'
    save_embeddings(filename, model, vocab_dict)
    
       

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Learning graph embeddings with path2vec')
    parser.add_argument('--input_file', required=True,
                        help='tab-separated gzipped file with training pairs and their similarities')
    parser.add_argument('--vsize', type=int, default=300, help='vector size')
    parser.add_argument('--bsize', type=int, default=100, help='batch size')
    parser.add_argument('--lrate', type=float, default=0.001, help='learning rate')
    parser.add_argument('--vocab_file', help='[optional] gzipped JSON file with the vocabulary (list of words)')
    # If the vocabulary file is not provided, it will be inferred from the training set
    # (can be painfully slow for large datasets)
    parser.add_argument('--fix_seeds', type=bool, default=False, help='fix seeds to ensure repeatability')
    parser.add_argument('--use_neighbors', type=bool, default=False,
                        help='whether or not to use the neighbor nodes-based regularizer')
    parser.add_argument('--neighbor_count', type=int, default=3,
                        help='number of adjacent nodes to consider for regularization')
    parser.add_argument('--negative_count', type=int, default=3, help='number of negative samples')
    parser.add_argument('--epochs', type=int, default=10, help='number of training epochs')
    parser.add_argument('--regularize', type=bool, default=False, help='L1 regularization of embeddings')
    parser.add_argument('--name', default='graph_emb', help='Run name, to be used in the file name')
    parser.add_argument('--l1factor', type=float, default=1e-10, help='L1 regularizer coefficient')
    parser.add_argument('--beta', type=float, default=0.01, help='neighbors-based regularizer first coefficient')
    parser.add_argument('--gamma', type=float, default=0.01, help='neighbors-based regularizer second coefficient')
    args = parser.parse_args()
    
    trainfile = args.input_file  # Gzipped file with pairs and their similarities
    embedding_dimension = args.vsize  # vector size (for example, 20)
    batch_size = args.bsize      # number of pairs in a batch (for example, 10)
    learn_rate = args.lrate   # Learning rate
    neighbors_count = args.neighbor_count
    negative = args.negative_count
    run_name = args.name
    l1_factor = args.l1factor
    beta = args.beta
    gamma = args.gamma
    
    print('Using adjacent nodes regularization: ', args.use_neighbors)
    
    run(trainfile = args.input_file,
        vocab_file = args.vocab_file,
        embedding_dimension = args.vsize,
        batch_size = args.bsize,
        learn_rate = args.lrate,
        neighbors_count = args.neighbor_count,
        negative = args.negative_count,
        run_name = args.name,
        l1_factor = args.l1factor,
        beta = args.beta,
        gamma = args.gamma,
        fix_seeds = args.fix_seeds,
        epochs = args.epochs,
        use_neighbors = args.use_neighbors,
        regularize = args.regularize)