run.py

import argparse
import time
from pprint import pprint
import numpy as np
import random
from pathlib import Path
import torch
from torch.utils.data import DataLoader
import dgl
from dgl.contrib.data import load_data
import os
import sys
import logging

from model import RGAT_DistMult, RGAT_ConvE, RGCN_ConvE, ConvE
from utils import process, TrainDataset, TestDataset, load_link


os.environ["TOKENIZERS_PARALLELISM"] = "false"

class Runner(object):
    def __init__(self, params):
        self.p = params
        self.prj_path = Path(__file__).parent.resolve()
        self.data = load_link(self.p.dataset)
        self.num_ent, self.train_data, self.valid_data, self.test_data, self.num_rels = self.data.num_nodes, self.data.train, self.data.valid, self.data.test, self.data.num_rels
        self.ent2textvector, self.rel2textvector, self.attr2vector = self.data.ent2textvector, self.data.rel2textvector,  self.data.attr2vector
        self.attrname, self.ent2value, self.ent2attrlabel = self.data.attrname, self.data.ent2value, self.data.ent2attrlabel
        self.ent2imgvector = self.data.ent2imgvector
        self.triplets = process({'train': self.train_data, 'valid': self.valid_data, 'test': self.test_data},
                                self.num_rels)
        if self.p.gpu != -1 and torch.cuda.is_available():
            self.device = torch.device(f'cuda:{self.p.gpu}')
        else:
            self.device = torch.device('cpu')
        self.p.embed_dim = self.p.k_w * self.p.k_h if self.p.embed_dim is None else self.p.embed_dim  # output dim of gnn
        self.data_iter = self.get_data_iter()
        self.g = self.build_graph()
        self.edge_type, self.edge_norm = self.get_edge_dir_and_norm()
        self.model = self.get_model()
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.p.lr, weight_decay=self.p.l2)
        self.best_val_mrr, self.best_epoch, self.best_val_results = 0., 0., {}
        self.logger = logging.getLogger()
        self.log_name = self.p.name + '.log'
        self.log_path = '/home/liangshuang/NewWork/logs'
        pprint(vars(self.p))

    def fit(self):
        epoch_hits1 = []
        epoch_hits3 = []
        epoch_hits10 = []
        epoch_mrr = []
        epoch_mr = []
        save_root = self.prj_path / 'checkpoints'
        if not save_root.exists():
            save_root.mkdir()
        save_path = save_root / (self.p.name + '.pt')

        if self.p.restore:
            self.load_model(save_path)
            print('Successfully Loaded previous model')

        for epoch in range(self.p.max_epochs):
            start_time = time.time()
            train_loss = self.train()
            val_results = self.evaluate('test')
            epoch_mr.append(val_results['mr'])
            epoch_mrr.append(val_results['mrr'])
            epoch_hits1.append(val_results['hits@1'])
            epoch_hits3.append(val_results['hits@3'])
            epoch_hits10.append(val_results['hits@10'])
            if val_results['mrr'] > self.best_val_mrr:
                self.best_val_results = val_results
                self.best_val_mrr = val_results['mrr']
                self.best_epoch = epoch
                self.save_model(save_path)
            print(
                f"[Epoch {epoch}]: Training Loss: {train_loss:.5}, Valid MRR: {val_results['mrr']:.5}, Best Valid MRR: {self.best_val_mrr:.5}, Cost: {time.time() - start_time:.2f}s")
        pprint(vars(self.p))
        self.load_model(save_path)
        print(f'Loading best model in {self.best_epoch} epoch, Evaluating on Test data')
        self.model.eval()
        # entity_embedding, rel_embedding = self.model.get_embedding(self.g)
        # np.save(os.path.join('./embedding', self.p.name + '_entity_embedding'), entity_embedding.cpu().numpy())
        # np.save(os.path.join('./embedding', self.p.name + '_rel_embedding'), rel_embedding.cpu().numpy())
        test_results = self.evaluate('test')
        print(
            f"MRR: Tail {test_results['left_mrr']:.5}, Head {test_results['right_mrr']:.5}, Avg {test_results['mrr']:.5}")
        print(f"MR: Tail {test_results['left_mr']:.5}, Head {test_results['right_mr']:.5}, Avg {test_results['mr']:.5}")
        print(f"hits@1 = {test_results['hits@1']:.5}")
        print(f"hits@3 = {test_results['hits@3']:.5}")
        print(f"hits@10 = {test_results['hits@10']:.5}")
        print("rank: ", self.model.rank)
        np.save(os.path.join('./result', self.p.name + '_epoch_mr'), epoch_mr)
        np.save(os.path.join('./result', self.p.name + '_epoch_mrr'), epoch_mrr)
        np.save(os.path.join('./result', self.p.name + '_epoch_hits@1'), epoch_hits1)
        np.save(os.path.join('./result', self.p.name + '_epoch_hits@3'), epoch_hits1)
        np.save(os.path.join('./result', self.p.name + '_epoch_hits@10'), epoch_hits1)


    def train(self):
        self.model.train()
        losses = []
        train_iter = self.data_iter['train']
        for step, (triplets, labels) in enumerate(train_iter):
            triplets, labels = triplets.to(self.device), labels.to(self.device)
            subj, rel = triplets[:, 0], triplets[:, 1]
            pred = self.model(self.g, subj, rel)  # [batch_size, num_ent]
            loss = self.model.calc_loss(pred, labels)

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            losses.append(loss.item())

        loss = np.mean(losses)
        return loss

    def evaluate(self, split):
        """
        Function to evaluate the model on validation or test set
        :param split: valid or test, set which data-set to evaluate on
        :return: results['mr']: Average of ranks_left and ranks_right
                 results['mrr']: Mean Reciprocal Rank
                 results['hits@k']: Probability of getting the correct prediction in top-k ranks based on predicted score
                 results['left_mrr'], results['left_mr'], results['right_mrr'], results['right_mr']
                 results['left_hits@k'], results['right_hits@k']
        """

        def get_combined_results(left, right):
            results = dict()
            assert left['count'] == right['count']
            count = float(left['count'])
            results['left_mr'] = round(left['mr'] / count, 5)
            results['left_mrr'] = round(left['mrr'] / count, 5)
            results['right_mr'] = round(right['mr'] / count, 5)
            results['right_mrr'] = round(right['mrr'] / count, 5)
            results['mr'] = round((left['mr'] + right['mr']) / (2 * count), 5)
            results['mrr'] = round((left['mrr'] + right['mrr']) / (2 * count), 5)
            for k in [1, 3, 10]:
                results[f'left_hits@{k}'] = round(left[f'hits@{k}'] / count, 5)
                results[f'right_hits@{k}'] = round(right[f'hits@{k}'] / count, 5)
                results[f'hits@{k}'] = round((results[f'left_hits@{k}'] + results[f'right_hits@{k}']) / 2, 5)
            return results

        self.model.eval()
        left_result = self.predict(split, 'tail')
        right_result = self.predict(split, 'head')
        res = get_combined_results(left_result, right_result)
        return res

    def predict(self, split='valid', mode='tail'):
        """
        Function to run model evaluation for a given mode
        :param split: valid or test, set which data-set to evaluate on
        :param mode: head or tail
        :return: results['mr']: Sum of ranks
                 results['mrr']: Sum of Reciprocal Rank
                 results['hits@k']: counts of getting the correct prediction in top-k ranks based on predicted score
                 results['count']: number of total predictions
        """
        with torch.no_grad():
            results = dict()
            test_iter = self.data_iter[f'{split}_{mode}']
            for step, (triplets, labels) in enumerate(test_iter):
                triplets, labels = triplets.to(self.device), labels.to(self.device)
                subj, rel, obj = triplets[:, 0], triplets[:, 1], triplets[:, 2]
                pred = self.model(self.g, subj, rel)
                b_range = torch.arange(pred.shape[0], device=self.device)
                target_pred = pred[b_range, obj]  # [batch_size, 1], get the predictive score of obj
                # label=>-1000000, not label=>pred, filter out other objects with same sub&rel pair
                pred = torch.where(labels.byte(), -torch.ones_like(pred) * 10000000, pred)
                pred[b_range, obj] = target_pred  # copy predictive score of obj to new pred
                ranks = 1 + torch.argsort(torch.argsort(pred, dim=1, descending=True), dim=1, descending=False)[
                    b_range, obj]  # get the rank of each (sub, rel, obj)
                ranks = ranks.float()
                results['count'] = torch.numel(ranks) + results.get('count', 0)  # number of predictions
                results['mr'] = torch.sum(ranks).item() + results.get('mr', 0)
                results['mrr'] = torch.sum(1.0 / ranks).item() + results.get('mrr', 0)

                for k in [1, 3, 10]:
                    results[f'hits@{k}'] = torch.numel(ranks[ranks <= k]) + results.get(f'hits@{k}', 0)
        return results

    def save_model(self, path):
        """
        Function to save a model. It saves the model parameters, best validation scores,
        best epoch corresponding to best validation, state of the optimizer and all arguments for the run.
        :param path: path where the model is saved
        :return:
        """
        state = {
            'model': self.model.state_dict(),
            'best_val': self.best_val_results,
            'best_epoch': self.best_epoch,
            'optimizer': self.optimizer.state_dict(),
            'args': vars(self.p)
        }
        torch.save(state, path)

    def load_model(self, path):
        """
        Function to load a saved model
        :param path: path where model is loaded
        :return:
        """
        state = torch.load(path)
        self.best_val_results = state['best_val']
        self.best_val_mrr = self.best_val_results['mrr']
        self.best_epoch = state['best_epoch']
        self.model.load_state_dict(state['model'])
        self.optimizer.load_state_dict(state['optimizer'])

    def build_graph(self):
        g = dgl.DGLGraph()
        g.add_nodes(self.num_ent)
        g.add_edges(self.train_data[:, 0], self.train_data[:, 2])
        g.add_edges(self.train_data[:, 2], self.train_data[:, 0])
        return g

    def get_data_iter(self):
        """
        get data loader for train, valid and test section
        :return: dict
        """

        def get_data_loader(dataset_class, split):
            return DataLoader(
                dataset_class(self.triplets[split], self.num_ent, self.p),
                batch_size=self.p.batch_size,
                shuffle=True,
                num_workers=self.p.num_workers
            )

        return {
            'train': get_data_loader(TrainDataset, 'train'),
            'valid_head': get_data_loader(TestDataset, 'valid_head'),
            'valid_tail': get_data_loader(TestDataset, 'valid_tail'),
            'test_head': get_data_loader(TestDataset, 'test_head'),
            'test_tail': get_data_loader(TestDataset, 'test_tail')
        }

    def get_edge_dir_and_norm(self):
        """
        :return: edge_type: indicates type of each edge: [E]
        """
        in_deg = self.g.in_degrees(range(self.g.number_of_nodes())).float().numpy()
        norm = in_deg ** -0.5
        norm[np.isinf(norm)] = 0
        self.g.ndata['xxx'] = norm
        self.g.apply_edges(lambda edges: {'xxx': edges.dst['xxx'] * edges.src['xxx']})
        norm = self.g.edata.pop('xxx').squeeze().to(self.device)
        edge_type = torch.tensor(np.concatenate([self.train_data[:, 1], self.train_data[:, 1] + self.num_rels])).to(
            self.device)
        return edge_type, norm

    def get_model(self):
        if self.p.model_func.lower() == 'distmult':
            model = RGAT_DistMult(num_ent=self.num_ent, num_rel=self.num_rels, num_base=self.p.num_bases,
                                     init_dim=self.p.init_dim, gcn_dim=self.p.gcn_dim, embed_dim=self.p.embed_dim,
                                     n_layer=self.p.n_layer, edge_type=self.edge_type, edge_norm=self.edge_norm,
                                     ent2textvector = self.ent2textvector, rel2textvector = self.rel2textvector,
                                     ent2attr = self.ent2value, ent2attrlabel = self.ent2attrlabel, attr2vector = self.attr2vector, ent2imgvector = self.ent2imgvector,
                                     use_text = self.p.text, use_img = self.p.img, use_attr = self.p.attr,
                                     device = self.device, bias=self.p.bias, gcn_drop=self.p.gcn_drop, opn=self.p.opn,
                                     hid_drop=self.p.hid_drop)
        elif self.p.model_func.lower() == 'conve':
            model = RGAT_ConvE(num_ent=self.num_ent, num_rel=self.num_rels, num_base=self.p.num_bases,
                                  init_dim=self.p.init_dim, gcn_dim=self.p.gcn_dim, embed_dim=self.p.embed_dim,
                                  n_layer=self.p.n_layer, edge_type=self.edge_type, edge_norm=self.edge_norm,
                                  ent2textvector=self.ent2textvector, rel2textvector=self.rel2textvector,
                                  ent2attr=self.ent2value, ent2attrlabel=self.ent2attrlabel, attr2vector=self.attr2vector, ent2imgvector = self.ent2imgvector,
                                  use_text=self.p.text, use_img=self.p.img, use_attr=self.p.attr,
                                  device=self.device,bias=self.p.bias, gcn_drop=self.p.gcn_drop, opn=self.p.opn,
                                  hid_drop=self.p.hid_drop, input_drop=self.p.input_drop,
                                  conve_hid_drop=self.p.conve_hid_drop, feat_drop=self.p.feat_drop,
                                  num_filt=self.p.num_filt, ker_sz=self.p.ker_sz, k_h=self.p.k_h, k_w=self.p.k_w)
        elif self.p.model_func.lower() == 'rgcn':
            model = RGCN_ConvE(num_ent=self.num_ent, num_rel=self.num_rels, num_base=self.p.num_bases,
                                  init_dim=self.p.init_dim, gcn_dim=self.p.gcn_dim, embed_dim=self.p.embed_dim,
                                  n_layer=self.p.n_layer, edge_type=self.edge_type, edge_norm=self.edge_norm,
                                  ent2textvector=self.ent2textvector, rel2textvector=self.rel2textvector,
                                  ent2attr=self.ent2value, ent2attrlabel=self.ent2attrlabel, attr2vector=self.attr2vector, ent2imgvector = self.ent2imgvector,
                                  use_text=self.p.text, use_img=self.p.img, use_attr=self.p.attr,
                                  device=self.device,bias=self.p.bias, gcn_drop=self.p.gcn_drop, opn=self.p.opn,
                                  hid_drop=self.p.hid_drop, input_drop=self.p.input_drop,
                                  conve_hid_drop=self.p.conve_hid_drop, feat_drop=self.p.feat_drop,
                                  num_filt=self.p.num_filt, ker_sz=self.p.ker_sz, k_h=self.p.k_h, k_w=self.p.k_w)
        elif self.p.model_func.lower() == 'simpleconve':
            model = ConvE(num_ent=self.num_ent, num_rel=self.num_rels, num_base=self.p.num_bases,
                                  init_dim=self.p.init_dim, gcn_dim=self.p.gcn_dim, embed_dim=self.p.embed_dim,
                                  n_layer=self.p.n_layer, edge_type=self.edge_type, edge_norm=self.edge_norm,
                                  ent2textvector=self.ent2textvector, rel2textvector=self.rel2textvector,
                                  ent2attr=self.ent2value, ent2attrlabel=self.ent2attrlabel, attr2vector=self.attr2vector, ent2imgvector = self.ent2imgvector,
                                  use_text=self.p.text, use_img=self.p.img, use_attr=self.p.attr,
                                  device=self.device,bias=self.p.bias, gcn_drop=self.p.gcn_drop, opn=self.p.opn,
                                  hid_drop=self.p.hid_drop, input_drop=self.p.input_drop,
                                  conve_hid_drop=self.p.conve_hid_drop, feat_drop=self.p.feat_drop,
                                  num_filt=self.p.num_filt, ker_sz=self.p.ker_sz, k_h=self.p.k_h, k_w=self.p.k_w)
        else:
            raise KeyError(f'score function {self.p.model_func} not recognized.')
        model.to(self.device)
        return model


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Parser For Arguments',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--name', default='test_run', help='Set run name for saving/restoring models')
    parser.add_argument('--data', dest='dataset', default='FB15k-237', help='Dataset to use, default: FB15k-237')
    parser.add_argument('--model_func', dest='model_func', default='conve',
                        help='Score Function for Link prediction')
    parser.add_argument('--opn', dest='opn', default='corr', help='Composition Operation to be used in CompGCN')

    parser.add_argument('--batch', dest='batch_size', default=256, type=int, help='Batch size')
    parser.add_argument('--gpu', type=int, default=0, help='Set GPU Ids : Eg: For CPU = -1, For Single GPU = 0')
    parser.add_argument('--epoch', dest='max_epochs', type=int, default=500, help='Number of epochs')
    parser.add_argument('--l2', type=float, default=0.0, help='L2 Regularization for Optimizer')
    parser.add_argument('--lr', type=float, default=0.001, help='Starting Learning Rate')
    parser.add_argument('--lbl_smooth', dest='lbl_smooth', type=float, default=0.1, help='Label Smoothing')
    parser.add_argument('--num_workers', type=int, default=8, help='Number of processes to construct batches')
    parser.add_argument('--seed', dest='seed', default=12345, type=int, help='Seed for randomization')

    parser.add_argument('--restore', dest='restore', action='store_true',
                        help='Restore from the previously saved model')
    parser.add_argument('--bias', dest='bias', action='store_true', help='Whether to use bias in the model')

    parser.add_argument('--text', dest='text', action='store_true', help='Whether to use text in the model')

    parser.add_argument('--img', dest='img',  action='store_true', help='Whether to use img in the model')

    parser.add_argument('--attr', dest='attr', action='store_true', help='Whether to use attr in the model')

    parser.add_argument('--num_bases', dest='num_bases', default=-1, type=int,
                        help='Number of basis relation vectors to use')
    parser.add_argument('--init_dim', dest='init_dim', default=100, type=int,
                        help='Initial dimension size for entities and relations')
    parser.add_argument('--gcn_dim', dest='gcn_dim', default=200, type=int, help='Number of hidden units in GCN')
    parser.add_argument('--embed_dim', dest='embed_dim', default=None, type=int,
                        help='Embedding dimension to give as input to score function')
    parser.add_argument('--n_layer', dest='n_layer', default=1, type=int, help='Number of GCN Layers to use')
    parser.add_argument('--gcn_drop', dest='gcn_drop', default=0.1, type=float, help='Dropout to use in GCN Layer')
    parser.add_argument('--hid_drop', dest='hid_drop', default=0.3, type=float, help='Dropout after GCN')

    # ConvE specific hyperparameters
    parser.add_argument('--conve_hid_drop', dest='conve_hid_drop', default=0.3, type=float,
                        help='ConvE: Hidden dropout')
    parser.add_argument('--feat_drop', dest='feat_drop', default=0.2, type=float, help='ConvE: Feature Dropout')
    parser.add_argument('--input_drop', dest='input_drop', default=0.2, type=float, help='ConvE: Stacked Input Dropout')
    parser.add_argument('--k_w', dest='k_w', default=20, type=int, help='ConvE: k_w')
    parser.add_argument('--k_h', dest='k_h', default=10, type=int, help='ConvE: k_h')
    parser.add_argument('--num_filt', dest='num_filt', default=200, type=int,
                        help='ConvE: Number of filters in convolution')
    parser.add_argument('--ker_sz', dest='ker_sz', default=7, type=int, help='ConvE: Kernel size to use')

    args = parser.parse_args()
    if not args.restore:
        args.name = time.strftime('%Y_%m_%d') + '_' + time.strftime(
            '%H:%M:%S') + '-' + args.model_func.lower() + '-' + args.opn

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    runner = Runner(args)
    runner.fit()