Sujit-O / pykg2vec / 759

pending completion

Build # 759

Build Type

Pull #177

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/Sujit-O/pykg2vec/commit/<a class=hub.com/Sujit-O/pykg2vec/commit/c981a3cd519666ff80051c69352510dfc3156ffa">c981a3cd5<a href="https://github.com/Sujit-O/pykg2vec/commit/c981a3cd519666ff80051c69352510dfc3156ffa">">Merge </a><a class="double-link" href="https://github.com/Sujit-O/pykg2vec/commit/<a class="double-link" href="https://github.com/Sujit-O/pykg2vec/commit/d7e59dcbc56f23241a9863053e00e0228010c0b2">d7e59dcbc</a>">d7e59dcbc</a><a href="https://github.com/Sujit-O/pykg2vec/commit/c981a3cd519666ff80051c69352510dfc3156ffa"> into 3fa2af7ad">3fa2af7ad</a>

Pull Request Pull Request #177: attempt to store hyper-parameters

Run Details

597 of 683 branches covered (87.41%)

Branch coverage included in aggregate %.

25 of 25 new or added lines in 2 files covered. (100.0%)

2892 of 3021 relevant lines covered (95.73%)

0.96 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

89.06

/pykg2vec/utils/trainer.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import warnings
import torch
import numpy as np
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd

from tqdm import tqdm
from pathlib import Path
from pykg2vec.utils.evaluator import Evaluator
from pykg2vec.utils.visualization import Visualization
from pykg2vec.data.generator import Generator
from pykg2vec.utils.logger import Logger
from pykg2vec.common import Monitor, TrainingStrategy
warnings.filterwarnings('ignore')


class EarlyStopper:

    """ Class used by trainer for handling the early stopping mechanism during the training of KGE algorithms.

        Args:
            patience (int): Number of epochs to wait before early stopping the training on no improvement.
            No early stopping if it is a negative number (default: {-1}).
            monitor (Monitor): the type of metric that earlystopper will monitor.

    """

    _logger = Logger().get_logger(__name__)

    def __init__(self, patience, monitor):

        self.monitor = monitor
        self.patience = patience

        # controlling variables.
        self.previous_metrics = None
        self.patience_left = patience

    def should_stop(self, curr_metrics):
        should_stop = False
        value, name = self.monitor.value, self.monitor.name

        if self.previous_metrics is not None:
            if self.monitor == Monitor.MEAN_RANK or self.monitor == Monitor.FILTERED_MEAN_RANK:
                is_worse = self.previous_metrics[value] < curr_metrics[value]
            else:
                is_worse = self.previous_metrics[value] > curr_metrics[value]

            if self.patience_left > 0 and is_worse:
                self.patience_left -= 1
                self._logger.info(
                    '%s more chances before the trainer stops the training. (prev_%s, curr_%s): (%.4f, %.4f)' %
                    (self.patience_left, name, name, self.previous_metrics[value], curr_metrics[value]))

            elif self.patience_left == 0 and is_worse:
                self._logger.info('Stop the training.')
                should_stop = True

            else:
                self._logger.info('Reset the patience count to %d' % (self.patience))
                self.patience_left = self.patience

        self.previous_metrics = curr_metrics

        return should_stop


class Trainer:
    """ Class for handling the training of the algorithms.

        Args:
            model (object): KGE model object

        Examples:
            >>> from pykg2vec.utils.trainer import Trainer
            >>> from pykg2vec.models.TransE import TransE
            >>> trainer = Trainer(TransE())
            >>> trainer.build_model()
            >>> trainer.train_model()
    """
    TRAINED_MODEL_FILE_NAME = "model.vec.pt"
    TRAINED_MODEL_CONFIG_NAME = "config.npy"
    _logger = Logger().get_logger(__name__)

    def __init__(self, model, config):
        self.model = model
        self.config = config

        self.best_metric = None
        self.monitor = None

        self.training_results = []

        self.evaluator = None
        self.generator = None
        self.optimizer = None
        self.early_stopper = None

    def build_model(self, monitor=Monitor.FILTERED_MEAN_RANK):
        """function to build the model"""
        if self.config.load_from_data is not None:
            self.load_model(self.config.load_from_data)

        self.evaluator = Evaluator(self.model, self.config)

        self.model.to(self.config.device)

        if self.config.optimizer == "adam":
            self.optimizer = optim.Adam(
                self.model.parameters(),
                lr=self.config.learning_rate,
            )
        elif self.config.optimizer == "sgd":
            self.optimizer = optim.SGD(
                self.model.parameters(),
                lr=self.config.learning_rate,
            )
        elif self.config.optimizer == "adagrad":
            self.optimizer = optim.Adagrad(
                self.model.parameters(),
                lr=self.config.learning_rate,
            )
        elif self.config.optimizer == "rms":
            self.optimizer = optim.RMSprop(
                self.model.parameters(),
                lr=self.config.learning_rate,
            )
        else:
            raise NotImplementedError("No support for %s optimizer" % self.config.optimizer)

        self.config.summary()

        self.early_stopper = EarlyStopper(self.config.patience, monitor)

    # Training related functions:
    def train_step_pairwise(self, pos_h, pos_r, pos_t, neg_h, neg_r, neg_t):
        pos_preds = self.model(pos_h, pos_r, pos_t)
        neg_preds = self.model(neg_h, neg_r, neg_t)

        if self.config.sampling == 'adversarial_negative_sampling':
            # RotatE: Adversarial Negative Sampling and alpha is the temperature.
            pos_preds = -pos_preds
            neg_preds = -neg_preds
            pos_preds = F.logsigmoid(pos_preds)
            neg_preds = neg_preds.view((-1, self.config.neg_rate))
            softmax = nn.Softmax(dim=1)(neg_preds*self.config.alpha).detach()
            neg_preds = torch.sum(softmax * (F.logsigmoid(-neg_preds)), dim=-1)
            loss = -neg_preds.mean() - pos_preds.mean()
        else:
            # others that use margin-based & pairwise loss function. (uniform or bern)
            loss = pos_preds + self.config.margin - neg_preds
            loss = torch.max(loss, torch.zeros_like(loss)).sum()

        if hasattr(self.model, 'get_reg'):
            # now only NTN uses regularizer,
            # other pairwise based KGE methods use normalization to regularize parameters.
            loss += self.model.get_reg()

        return loss

    def train_step_projection(self, h, r, t, hr_t, tr_h):
        if self.model.model_name.lower() == "conve" or self.model.model_name.lower() == "tucker":
            if hasattr(self.config, 'label_smoothing'):
                hr_t = hr_t * (1.0 - self.config.label_smoothing) + 1.0 / self.config.tot_entity
                tr_h = tr_h * (1.0 - self.config.label_smoothing) + 1.0 / self.config.tot_entity

            pred_tails = self.model(h, r, direction="tail")  # (h, r) -> hr_t forward
            pred_heads = self.model(t, r, direction="head")  # (t, r) -> tr_h backward

            loss_tails = torch.mean(F.binary_cross_entropy(pred_tails, hr_t))
            loss_heads = torch.mean(F.binary_cross_entropy(pred_heads, tr_h))

            loss = loss_tails + loss_heads

        else:
            loss_tails = self.model(h, r, hr_t, direction="tail")  # (h, r) -> hr_t forward
            loss_heads = self.model(t, r, tr_h, direction="head")  # (t, r) -> tr_h backward

            loss = loss_tails + loss_heads

            if hasattr(self.model, 'get_reg'):
                # now only complex distmult uses regularizer in algorithms,
                loss += self.model.get_reg()

        return loss

    def train_step_pointwise(self, h, r, t, y):
        preds = self.model(h, r, t)
        loss = F.softplus(y*preds).mean()

        if hasattr(self.model, 'get_reg'): # for complex & complex-N3 & DistMult & CP & ANALOGY
            loss += self.model.get_reg(h, r, t)

        return loss

    def train_model(self):

        # for key, value in self.config.__dict__.items():
        #     print(key," ",value)
        #print(self.config.__dict__[""])
        # pdb.set_trace()

        """Function to train the model."""
        self.generator = Generator(self.model, self.config)
        self.monitor = Monitor.FILTERED_MEAN_RANK
        for cur_epoch_idx in range(self.config.epochs):
            self._logger.info("Epoch[%d/%d]" % (cur_epoch_idx, self.config.epochs))

            self.train_model_epoch(cur_epoch_idx)

            if cur_epoch_idx % self.config.test_step == 0:
                self.model.eval()
                metrics = self.evaluator.mini_test(cur_epoch_idx)

                if self.early_stopper.should_stop(metrics):
                    ### Early Stop Mechanism
                    ### start to check if the metric is still improving after each mini-test.
                    ### Example, if test_step == 5, the trainer will check metrics every 5 epoch.
                    break

                # store the best model weights.
                if self.config.save_model:
                    if self.best_metric is None:
                        self.best_metric = metrics
                        self.save_model()
                    else:
                        if self.monitor == Monitor.MEAN_RANK or self.monitor == Monitor.FILTERED_MEAN_RANK:
                            is_better = self.best_metric[self.monitor.value] > metrics[self.monitor.value]
                        else:
                            is_better = self.best_metric[self.monitor.value] < metrics[self.monitor.value]
                        if is_better:
                            self.save_model()
                            self.best_metric = metrics

        self.evaluator.full_test(cur_epoch_idx)
        self.evaluator.metric_calculator.save_test_summary(self.model.model_name)

        self.generator.stop()
        self.save_training_result()

        # if self.config.save_model:
        #     self.save_model()

        if self.config.disp_result:
            self.display()

        self.export_embeddings()

        return cur_epoch_idx # the runned epoches.

    def tune_model(self):
        """Function to tune the model."""
        current_loss = float("inf")

        self.generator = Generator(self.model, self.config)
        self.evaluator = Evaluator(self.model, self.config, tuning=True)

        for cur_epoch_idx in range(self.config.epochs):
            current_loss = self.train_model_epoch(cur_epoch_idx, tuning=True)

        self.evaluator.full_test(cur_epoch_idx)

        self.generator.stop()

        return current_loss

    def train_model_epoch(self, epoch_idx, tuning=False):
        """Function to train the model for one epoch."""
        acc_loss = 0

        num_batch = self.config.tot_train_triples // self.config.batch_size if not self.config.debug else 10

        self.generator.start_one_epoch(num_batch)

        progress_bar = tqdm(range(num_batch))

        for _ in progress_bar:
            data = list(next(self.generator))

            self.model.train()
            self.optimizer.zero_grad()

            if self.model.training_strategy == TrainingStrategy.PROJECTION_BASED:
                h = torch.LongTensor(data[0]).to(self.config.device)
                r = torch.LongTensor(data[1]).to(self.config.device)
                t = torch.LongTensor(data[2]).to(self.config.device)
                hr_t = data[3].to(self.config.device)
                tr_h = data[4].to(self.config.device)
                loss = self.train_step_projection(h, r, t, hr_t, tr_h)
            elif self.model.training_strategy == TrainingStrategy.POINTWISE_BASED:
                h = torch.LongTensor(data[0]).to(self.config.device)
                r = torch.LongTensor(data[1]).to(self.config.device)
                t = torch.LongTensor(data[2]).to(self.config.device)
                y = torch.LongTensor(data[3]).to(self.config.device)
                loss = self.train_step_pointwise(h, r, t, y)
            elif self.model.training_strategy == TrainingStrategy.PAIRWISE_BASED:
                pos_h = torch.LongTensor(data[0]).to(self.config.device)
                pos_r = torch.LongTensor(data[1]).to(self.config.device)
                pos_t = torch.LongTensor(data[2]).to(self.config.device)
                neg_h = torch.LongTensor(data[3]).to(self.config.device)
                neg_r = torch.LongTensor(data[4]).to(self.config.device)
                neg_t = torch.LongTensor(data[5]).to(self.config.device)
                loss = self.train_step_pairwise(pos_h, pos_r, pos_t, neg_h, neg_r, neg_t)
            else:
                raise NotImplementedError("Unknown training strategy: %s" % self.model.training_strategy)

            loss.backward()
            self.optimizer.step()
            acc_loss += loss.item()

            if not tuning:
                progress_bar.set_description('acc_loss: %f, cur_loss: %f'% (acc_loss, loss))

        self.training_results.append([epoch_idx, acc_loss])

        return acc_loss

    def enter_interactive_mode(self):
        self.build_model()
        self.load_model()

        self._logger.info("""The training/loading of the model has finished!
                                    Now enter interactive mode :)
                                    -----
                                    Example 1: trainer.infer_tails(1,10,topk=5)""")
        self.infer_tails(1, 10, topk=5)

        self._logger.info("""-----
                                    Example 2: trainer.infer_heads(10,20,topk=5)""")
        self.infer_heads(10, 20, topk=5)

        self._logger.info("""-----
                                    Example 3: trainer.infer_rels(1,20,topk=5)""")
        self.infer_rels(1, 20, topk=5)

    def exit_interactive_mode(self):
        self._logger.info("Thank you for trying out inference interactive script :)")

    def infer_tails(self, h, r, topk=5):
        tails = self.evaluator.test_tail_rank(h, r, topk).cpu().numpy()
        idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')
        idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')
        logs = [
            "",
            "(head, relation)->({},{}) :: Inferred tails->({})".format(h, r, ",".join([str(i) for i in tails])),
            "",
            "head: %s" % idx2ent[h],
            "relation: %s" % idx2rel[r],
        ]

        for idx, tail in enumerate(tails):
            logs.append("%dth predicted tail: %s" % (idx, idx2ent[tail]))

        self._logger.info("\n".join(logs))
        return {tail: idx2ent[tail] for tail in tails}

    def infer_heads(self, r, t, topk=5):
        heads = self.evaluator.test_head_rank(r, t, topk).cpu().numpy()
        idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')
        idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')
        logs = [
            "",
            "(relation,tail)->({},{}) :: Inferred heads->({})".format(t, r, ",".join([str(i) for i in heads])),
            "",
            "tail: %s" % idx2ent[t],
            "relation: %s" % idx2rel[r],
        ]

        for idx, head in enumerate(heads):
            logs.append("%dth predicted head: %s" % (idx, idx2ent[head]))

        self._logger.info("\n".join(logs))
        return {head: idx2ent[head] for head in heads}

    def infer_rels(self, h, t, topk=5):
        if self.model.model_name.lower() in ["proje_pointwise", "conve", "tucker"]:
            self._logger.info("%s model doesn't support relation inference in nature.")
            return {}

        rels = self.evaluator.test_rel_rank(h, t, topk).cpu().numpy()
        idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')
        idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')
        logs = [
            "",
            "(head,tail)->({},{}) :: Inferred rels->({})".format(h, t, ",".join([str(i) for i in rels])),
            "",
            "head: %s" % idx2ent[h],
            "tail: %s" % idx2ent[t],
        ]

        for idx, rel in enumerate(rels):
            logs.append("%dth predicted rel: %s" % (idx, idx2rel[rel]))

        self._logger.info("\n".join(logs))
        return {rel: idx2rel[rel] for rel in rels}

    # ''' Procedural functions:'''
    def save_model(self):
        """Function to save the model."""
        saved_path = self.config.path_tmp / self.model.model_name
        saved_path.mkdir(parents=True, exist_ok=True)
        torch.save(self.model.state_dict(), str(saved_path / self.TRAINED_MODEL_FILE_NAME))

        """Save hyper-parameters into a yaml file with the model"""
        save_path_config = saved_path / self.TRAINED_MODEL_CONFIG_NAME
        np.save(save_path_config, self.config)

    def load_model(self, model_path=None):
        """Function to load the model."""
        if model_path is None:
            model_path = self.config.path_tmp / self.model.model_name / self.TRAINED_MODEL_FILE_NAME
            model_path_file = self.config.path_tmp / self.model.model_name / self.TRAINED_MODEL_FILE_NAME
            model_path_config = self.config.path_tmp / self.model.model_name / self.TRAINED_MODEL_CONFIG_NAME
        else:
            model_path = Path(model_path)
            model_path_file = model_path / self.TRAINED_MODEL_FILE_NAME
            model_path_config = model_path / self.TRAINED_MODEL_CONFIG_NAME

        if model_path_file.exists() and model_path_config.exists():
            config_temp = np.load(model_path_config, allow_pickle=True).item()
            # for key, value in config_temp.__dict__.items():
            #     print(key, " ", value)
            self.config.__dict__['lmbda'] = config_temp.__dict__['lmbda']
            self.config.__dict__['l1_flag'] = config_temp.__dict__['l1_flag']
            self.config.__dict__['learning_rate'] = config_temp.__dict__['learning_rate']
            self.config.__dict__['hidden_size'] = config_temp.__dict__['hidden_size']
            self.config.__dict__['batch_size'] = config_temp.__dict__['batch_size']
            self.config.__dict__['epochs'] = config_temp.__dict__['epochs']
            self.config.__dict__['margin'] = config_temp.__dict__['margin']
            self.config.__dict__['optimizer'] = config_temp.__dict__['optimizer']
            self.config.__dict__['sampling'] = config_temp.__dict__['sampling']
            self.config.__dict__['neg_rate'] = config_temp.__dict__['neg_rate']

            self.model.__init__(**self.config.__dict__)

            self.model.load_state_dict(torch.load(str(model_path_file)))
            self.model.eval()
        else:
            raise ValueError("Cannot load model from %s" % model_path_file)

    def display(self):
        """Function to display embedding."""
        options = {"ent_only_plot": True,
                   "rel_only_plot": not self.config.plot_entity_only,
                   "ent_and_rel_plot": not self.config.plot_entity_only}

        if self.config.plot_embedding:
            viz = Visualization(self.model, self.config, vis_opts=options)
            viz.plot_embedding(resultpath=self.config.path_figures, algos=self.model.model_name, show_label=False)

        if self.config.plot_training_result:
            viz = Visualization(self.model, self.config)
            viz.plot_train_result()

        if self.config.plot_testing_result:
            viz = Visualization(self.model, self.config)
            viz.plot_test_result()

    def export_embeddings(self):
        """
            Export embeddings in tsv and pandas pickled format.
            With tsvs (both label, vector files), you can:
            1) Use those pretained embeddings for your applications.
            2) Visualize the embeddings in this website to gain insights. (https://projector.tensorflow.org/)

            Pandas dataframes can be read with pd.read_pickle('desired_file.pickle')
        """
        save_path = self.config.path_embeddings / self.model.model_name
        save_path.mkdir(parents=True, exist_ok=True)

        idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')
        idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')

        with open(str(save_path / "ent_labels.tsv"), 'w') as l_export_file:
            for label in idx2ent.values():
                l_export_file.write(label + "\n")

        with open(str(save_path / "rel_labels.tsv"), 'w') as l_export_file:
            for label in idx2rel.values():
                l_export_file.write(label + "\n")

        for named_embedding in self.model.parameter_list:
            all_ids = list(range(0, int(named_embedding.weight.shape[0])))

            stored_name = named_embedding.name

            if len(named_embedding.weight.shape) == 2:
                all_embs = named_embedding.weight.detach().cpu().numpy()
                with open(str(save_path / ("%s.tsv" % stored_name)), 'w') as v_export_file:
                    for idx in all_ids:
                        v_export_file.write("\t".join([str(x) for x in all_embs[idx]]) + "\n")

    def save_training_result(self):
        """Function that saves training result"""
        files = os.listdir(str(self.config.path_result))
        l = len([f for f in files if self.model.model_name in f if 'Training' in f])
        df = pd.DataFrame(self.training_results, columns=['Epochs', 'Loss'])
        with open(str(self.config.path_result / (self.model.model_name + '_Training_results_' + str(l) + '.csv')),
                  'w') as fh:
            df.to_csv(fh)

1	#!/usr/bin/env python
2	# -- coding: utf-8 --
3	import os	1✔
4	import warnings	1✔
5	import torch	1✔
6	import numpy as np	1✔
7	import torch.optim as optim	1✔
8	import torch.nn as nn	1✔
9	import torch.nn.functional as F	1✔
10	import pandas as pd	1✔
11
12	from tqdm import tqdm	1✔
13	from pathlib import Path	1✔
14	from pykg2vec.utils.evaluator import Evaluator	1✔
15	from pykg2vec.utils.visualization import Visualization	1✔
16	from pykg2vec.data.generator import Generator	1✔
17	from pykg2vec.utils.logger import Logger	1✔
18	from pykg2vec.common import Monitor, TrainingStrategy	1✔
19	warnings.filterwarnings('ignore')	1✔
20
21
22	class EarlyStopper:	1✔
23
24	""" Class used by trainer for handling the early stopping mechanism during the training of KGE algorithms.
25
26	Args:
27	patience (int): Number of epochs to wait before early stopping the training on no improvement.
28	No early stopping if it is a negative number (default: {-1}).
29	monitor (Monitor): the type of metric that earlystopper will monitor.
30
31	"""
32
33	_logger = Logger().get_logger(__name__)	1✔
34
35	def __init__(self, patience, monitor):	1✔
36
37	self.monitor = monitor	1✔
38	self.patience = patience	1✔
39
40	# controlling variables.
41	self.previous_metrics = None	1✔
42	self.patience_left = patience	1✔
43
44	def should_stop(self, curr_metrics):	1✔
45	should_stop = False	1✔
46	value, name = self.monitor.value, self.monitor.name	1✔
47
48	if self.previous_metrics is not None:	1✔
49	if self.monitor == Monitor.MEAN_RANK or self.monitor == Monitor.FILTERED_MEAN_RANK:	1✔
50	is_worse = self.previous_metrics[value] < curr_metrics[value]	1✔
51	else:
52	is_worse = self.previous_metrics[value] > curr_metrics[value]	1✔
53
54	if self.patience_left > 0 and is_worse:	1!
55	self.patience_left -= 1	×
56	self._logger.info(	×
57	'%s more chances before the trainer stops the training. (prev_%s, curr_%s): (%.4f, %.4f)' %
58	(self.patience_left, name, name, self.previous_metrics[value], curr_metrics[value]))
59
60	elif self.patience_left == 0 and is_worse:	1✔
61	self._logger.info('Stop the training.')	1✔
62	should_stop = True	1✔
63
64	else:
65	self._logger.info('Reset the patience count to %d' % (self.patience))	1✔
66	self.patience_left = self.patience	1✔
67
68	self.previous_metrics = curr_metrics	1✔
69
70	return should_stop	1✔
71
72
73	class Trainer:	1✔
74	""" Class for handling the training of the algorithms.
75
76	Args:
77	model (object): KGE model object
78
79	Examples:
80	>>> from pykg2vec.utils.trainer import Trainer
81	>>> from pykg2vec.models.TransE import TransE
82	>>> trainer = Trainer(TransE())
83	>>> trainer.build_model()
84	>>> trainer.train_model()
85	"""
86	TRAINED_MODEL_FILE_NAME = "model.vec.pt"	1✔
87	TRAINED_MODEL_CONFIG_NAME = "config.npy"	1✔
88	_logger = Logger().get_logger(__name__)	1✔
89
90	def __init__(self, model, config):	1✔
91	self.model = model	1✔
92	self.config = config	1✔
93
94	self.best_metric = None	1✔
95	self.monitor = None	1✔
96
97	self.training_results = []	1✔
98
99	self.evaluator = None	1✔
100	self.generator = None	1✔
101	self.optimizer = None	1✔
102	self.early_stopper = None	1✔
103
104	def build_model(self, monitor=Monitor.FILTERED_MEAN_RANK):	1✔
105	"""function to build the model"""
106	if self.config.load_from_data is not None:	1!
107	self.load_model(self.config.load_from_data)	×
108
109	self.evaluator = Evaluator(self.model, self.config)	1✔
110
111	self.model.to(self.config.device)	1✔
112
113	if self.config.optimizer == "adam":	1✔
114	self.optimizer = optim.Adam(	1✔
115	self.model.parameters(),
116	lr=self.config.learning_rate,
117	)
118	elif self.config.optimizer == "sgd":	1✔
119	self.optimizer = optim.SGD(	1✔
120	self.model.parameters(),
121	lr=self.config.learning_rate,
122	)
123	elif self.config.optimizer == "adagrad":	1!
124	self.optimizer = optim.Adagrad(	×
125	self.model.parameters(),
126	lr=self.config.learning_rate,
127	)
128	elif self.config.optimizer == "rms":	1!
129	self.optimizer = optim.RMSprop(	1✔
130	self.model.parameters(),
131	lr=self.config.learning_rate,
132	)
133	else:
134	raise NotImplementedError("No support for %s optimizer" % self.config.optimizer)	×
135
136	self.config.summary()	1✔
137
138	self.early_stopper = EarlyStopper(self.config.patience, monitor)	1✔
139
140	# Training related functions:
141	def train_step_pairwise(self, pos_h, pos_r, pos_t, neg_h, neg_r, neg_t):	1✔
142	pos_preds = self.model(pos_h, pos_r, pos_t)	1✔
143	neg_preds = self.model(neg_h, neg_r, neg_t)	1✔
144
145	if self.config.sampling == 'adversarial_negative_sampling':	1!
146	# RotatE: Adversarial Negative Sampling and alpha is the temperature.
147	pos_preds = -pos_preds	×
148	neg_preds = -neg_preds	×
149	pos_preds = F.logsigmoid(pos_preds)	×
150	neg_preds = neg_preds.view((-1, self.config.neg_rate))	×
151	softmax = nn.Softmax(dim=1)(neg_preds*self.config.alpha).detach()	×
152	neg_preds = torch.sum(softmax * (F.logsigmoid(-neg_preds)), dim=-1)	×
153	loss = -neg_preds.mean() - pos_preds.mean()	×
154	else:
155	# others that use margin-based & pairwise loss function. (uniform or bern)
156	loss = pos_preds + self.config.margin - neg_preds	1✔
157	loss = torch.max(loss, torch.zeros_like(loss)).sum()	1✔
158
159	if hasattr(self.model, 'get_reg'):	1✔
160	# now only NTN uses regularizer,
161	# other pairwise based KGE methods use normalization to regularize parameters.
162	loss += self.model.get_reg()	1✔
163
164	return loss	1✔
165
166	def train_step_projection(self, h, r, t, hr_t, tr_h):	1✔
167	if self.model.model_name.lower() == "conve" or self.model.model_name.lower() == "tucker":	1✔
168	if hasattr(self.config, 'label_smoothing'):	1!
169	hr_t = hr_t * (1.0 - self.config.label_smoothing) + 1.0 / self.config.tot_entity	1✔
170	tr_h = tr_h * (1.0 - self.config.label_smoothing) + 1.0 / self.config.tot_entity	1✔
171
172	pred_tails = self.model(h, r, direction="tail") # (h, r) -> hr_t forward	1✔
173	pred_heads = self.model(t, r, direction="head") # (t, r) -> tr_h backward	1✔
174
175	loss_tails = torch.mean(F.binary_cross_entropy(pred_tails, hr_t))	1✔
176	loss_heads = torch.mean(F.binary_cross_entropy(pred_heads, tr_h))	1✔
177
178	loss = loss_tails + loss_heads	1✔
179
180	else:
181	loss_tails = self.model(h, r, hr_t, direction="tail") # (h, r) -> hr_t forward	1✔
182	loss_heads = self.model(t, r, tr_h, direction="head") # (t, r) -> tr_h backward	1✔
183
184	loss = loss_tails + loss_heads	1✔
185
186	if hasattr(self.model, 'get_reg'):	1!
187	# now only complex distmult uses regularizer in algorithms,
188	loss += self.model.get_reg()	1✔
189
190	return loss	1✔
191
192	def train_step_pointwise(self, h, r, t, y):	1✔
193	preds = self.model(h, r, t)	1✔
194	loss = F.softplus(y*preds).mean()	1✔
195
196	if hasattr(self.model, 'get_reg'): # for complex & complex-N3 & DistMult & CP & ANALOGY	1✔
197	loss += self.model.get_reg(h, r, t)	1✔
198
199	return loss	1✔
200
201	def train_model(self):	1✔
202
203	# for key, value in self.config.__dict__.items():
204	# print(key," ",value)
205	#print(self.config.__dict__[""])
206	# pdb.set_trace()
207
208	"""Function to train the model."""
209	self.generator = Generator(self.model, self.config)	1✔
210	self.monitor = Monitor.FILTERED_MEAN_RANK	1✔
211	for cur_epoch_idx in range(self.config.epochs):	1✔
212	self._logger.info("Epoch[%d/%d]" % (cur_epoch_idx, self.config.epochs))	1✔
213
214	self.train_model_epoch(cur_epoch_idx)	1✔
215
216	if cur_epoch_idx % self.config.test_step == 0:	1!
217	self.model.eval()	1✔
218	metrics = self.evaluator.mini_test(cur_epoch_idx)	1✔
219
220	if self.early_stopper.should_stop(metrics):	1✔
221	### Early Stop Mechanism
222	### start to check if the metric is still improving after each mini-test.
223	### Example, if test_step == 5, the trainer will check metrics every 5 epoch.
224	break	1✔
225
226	# store the best model weights.
227	if self.config.save_model:	1✔
228	if self.best_metric is None:	1!
229	self.best_metric = metrics	1✔
230	self.save_model()	1✔
231	else:
232	if self.monitor == Monitor.MEAN_RANK or self.monitor == Monitor.FILTERED_MEAN_RANK:	×
233	is_better = self.best_metric[self.monitor.value] > metrics[self.monitor.value]	×
234	else:
235	is_better = self.best_metric[self.monitor.value] < metrics[self.monitor.value]	×
236	if is_better:	×
237	self.save_model()	×
238	self.best_metric = metrics	×
239
240	self.evaluator.full_test(cur_epoch_idx)	1✔
241	self.evaluator.metric_calculator.save_test_summary(self.model.model_name)	1✔
242
243	self.generator.stop()	1✔
244	self.save_training_result()	1✔
245
246	# if self.config.save_model:
247	# self.save_model()
248
249	if self.config.disp_result:	1✔
250	self.display()	1✔
251
252	self.export_embeddings()	1✔
253
254	return cur_epoch_idx # the runned epoches.	1✔
255
256	def tune_model(self):	1✔
257	"""Function to tune the model."""
258	current_loss = float("inf")	1✔
259
260	self.generator = Generator(self.model, self.config)	1✔
261	self.evaluator = Evaluator(self.model, self.config, tuning=True)	1✔
262
263	for cur_epoch_idx in range(self.config.epochs):	1✔
264	current_loss = self.train_model_epoch(cur_epoch_idx, tuning=True)	1✔
265
266	self.evaluator.full_test(cur_epoch_idx)	1✔
267
268	self.generator.stop()	1✔
269
270	return current_loss	1✔
271
272	def train_model_epoch(self, epoch_idx, tuning=False):	1✔
273	"""Function to train the model for one epoch."""
274	acc_loss = 0	1✔
275
276	num_batch = self.config.tot_train_triples // self.config.batch_size if not self.config.debug else 10	1✔
277
278	self.generator.start_one_epoch(num_batch)	1✔
279
280	progress_bar = tqdm(range(num_batch))	1✔
281
282	for _ in progress_bar:	1✔
283	data = list(next(self.generator))	1✔
284
285	self.model.train()	1✔
286	self.optimizer.zero_grad()	1✔
287
288	if self.model.training_strategy == TrainingStrategy.PROJECTION_BASED:	1✔
289	h = torch.LongTensor(data[0]).to(self.config.device)	1✔
290	r = torch.LongTensor(data[1]).to(self.config.device)	1✔
291	t = torch.LongTensor(data[2]).to(self.config.device)	1✔
292	hr_t = data[3].to(self.config.device)	1✔
293	tr_h = data[4].to(self.config.device)	1✔
294	loss = self.train_step_projection(h, r, t, hr_t, tr_h)	1✔
295	elif self.model.training_strategy == TrainingStrategy.POINTWISE_BASED:	1✔
296	h = torch.LongTensor(data[0]).to(self.config.device)	1✔
297	r = torch.LongTensor(data[1]).to(self.config.device)	1✔
298	t = torch.LongTensor(data[2]).to(self.config.device)	1✔
299	y = torch.LongTensor(data[3]).to(self.config.device)	1✔
300	loss = self.train_step_pointwise(h, r, t, y)	1✔
301	elif self.model.training_strategy == TrainingStrategy.PAIRWISE_BASED:	1!
302	pos_h = torch.LongTensor(data[0]).to(self.config.device)	1✔
303	pos_r = torch.LongTensor(data[1]).to(self.config.device)	1✔
304	pos_t = torch.LongTensor(data[2]).to(self.config.device)	1✔
305	neg_h = torch.LongTensor(data[3]).to(self.config.device)	1✔
306	neg_r = torch.LongTensor(data[4]).to(self.config.device)	1✔
307	neg_t = torch.LongTensor(data[5]).to(self.config.device)	1✔
308	loss = self.train_step_pairwise(pos_h, pos_r, pos_t, neg_h, neg_r, neg_t)	1✔
309	else:
310	raise NotImplementedError("Unknown training strategy: %s" % self.model.training_strategy)	×
311
312	loss.backward()	1✔
313	self.optimizer.step()	1✔
314	acc_loss += loss.item()	1✔
315
316	if not tuning:	1✔
317	progress_bar.set_description('acc_loss: %f, cur_loss: %f'% (acc_loss, loss))	1✔
318
319	self.training_results.append([epoch_idx, acc_loss])	1✔
320
321	return acc_loss	1✔
322
323	def enter_interactive_mode(self):	1✔
324	self.build_model()	1✔
325	self.load_model()	1✔
326
327	self._logger.info("""The training/loading of the model has finished!	1✔
328	Now enter interactive mode :)
329	-----
330	Example 1: trainer.infer_tails(1,10,topk=5)""")
331	self.infer_tails(1, 10, topk=5)	1✔
332
333	self._logger.info("""-----	1✔
334	Example 2: trainer.infer_heads(10,20,topk=5)""")
335	self.infer_heads(10, 20, topk=5)	1✔
336
337	self._logger.info("""-----	1✔
338	Example 3: trainer.infer_rels(1,20,topk=5)""")
339	self.infer_rels(1, 20, topk=5)	1✔
340
341	def exit_interactive_mode(self):	1✔
342	self._logger.info("Thank you for trying out inference interactive script :)")	1✔
343
344	def infer_tails(self, h, r, topk=5):	1✔
345	tails = self.evaluator.test_tail_rank(h, r, topk).cpu().numpy()	1✔
346	idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')	1✔
347	idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')	1✔
348	logs = [	1✔
349	"",
350	"(head, relation)->({},{}) :: Inferred tails->({})".format(h, r, ",".join([str(i) for i in tails])),
351	"",
352	"head: %s" % idx2ent[h],
353	"relation: %s" % idx2rel[r],
354	]
355
356	for idx, tail in enumerate(tails):	1✔
357	logs.append("%dth predicted tail: %s" % (idx, idx2ent[tail]))	1✔
358
359	self._logger.info("\n".join(logs))	1✔
360	return {tail: idx2ent[tail] for tail in tails}	1✔
361
362	def infer_heads(self, r, t, topk=5):	1✔
363	heads = self.evaluator.test_head_rank(r, t, topk).cpu().numpy()	1✔
364	idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')	1✔
365	idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')	1✔
366	logs = [	1✔
367	"",
368	"(relation,tail)->({},{}) :: Inferred heads->({})".format(t, r, ",".join([str(i) for i in heads])),
369	"",
370	"tail: %s" % idx2ent[t],
371	"relation: %s" % idx2rel[r],
372	]
373
374	for idx, head in enumerate(heads):	1✔
375	logs.append("%dth predicted head: %s" % (idx, idx2ent[head]))	1✔
376
377	self._logger.info("\n".join(logs))	1✔
378	return {head: idx2ent[head] for head in heads}	1✔
379
380	def infer_rels(self, h, t, topk=5):	1✔
381	if self.model.model_name.lower() in ["proje_pointwise", "conve", "tucker"]:	1✔
382	self._logger.info("%s model doesn't support relation inference in nature.")	1✔
383	return {}	1✔
384
385	rels = self.evaluator.test_rel_rank(h, t, topk).cpu().numpy()	1✔
386	idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')	1✔
387	idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')	1✔
388	logs = [	1✔
389	"",
390	"(head,tail)->({},{}) :: Inferred rels->({})".format(h, t, ",".join([str(i) for i in rels])),
391	"",
392	"head: %s" % idx2ent[h],
393	"tail: %s" % idx2ent[t],
394	]
395
396	for idx, rel in enumerate(rels):	1✔
397	logs.append("%dth predicted rel: %s" % (idx, idx2rel[rel]))	1✔
398
399	self._logger.info("\n".join(logs))	1✔
400	return {rel: idx2rel[rel] for rel in rels}	1✔
401
402	# ''' Procedural functions:'''
403	def save_model(self):	1✔
404	"""Function to save the model."""
405	saved_path = self.config.path_tmp / self.model.model_name	1✔
406	saved_path.mkdir(parents=True, exist_ok=True)	1✔
407	torch.save(self.model.state_dict(), str(saved_path / self.TRAINED_MODEL_FILE_NAME))	1✔
408
409	"""Save hyper-parameters into a yaml file with the model"""
410	save_path_config = saved_path / self.TRAINED_MODEL_CONFIG_NAME	1✔
411	np.save(save_path_config, self.config)	1✔
412
413	def load_model(self, model_path=None):	1✔
414	"""Function to load the model."""
415	if model_path is None:	1!
416	model_path = self.config.path_tmp / self.model.model_name / self.TRAINED_MODEL_FILE_NAME	1✔
417	model_path_file = self.config.path_tmp / self.model.model_name / self.TRAINED_MODEL_FILE_NAME	1✔
418	model_path_config = self.config.path_tmp / self.model.model_name / self.TRAINED_MODEL_CONFIG_NAME	1✔
419	else:
420	model_path = Path(model_path)	×
421	model_path_file = model_path / self.TRAINED_MODEL_FILE_NAME	×
422	model_path_config = model_path / self.TRAINED_MODEL_CONFIG_NAME	×
423
424	if model_path_file.exists() and model_path_config.exists():	1!
425	config_temp = np.load(model_path_config, allow_pickle=True).item()	1✔
426	# for key, value in config_temp.__dict__.items():
427	# print(key, " ", value)
428	self.config.__dict__['lmbda'] = config_temp.__dict__['lmbda']	1✔
429	self.config.__dict__['l1_flag'] = config_temp.__dict__['l1_flag']	1✔
430	self.config.__dict__['learning_rate'] = config_temp.__dict__['learning_rate']	1✔
431	self.config.__dict__['hidden_size'] = config_temp.__dict__['hidden_size']	1✔
432	self.config.__dict__['batch_size'] = config_temp.__dict__['batch_size']	1✔
433	self.config.__dict__['epochs'] = config_temp.__dict__['epochs']	1✔
434	self.config.__dict__['margin'] = config_temp.__dict__['margin']	1✔
435	self.config.__dict__['optimizer'] = config_temp.__dict__['optimizer']	1✔
436	self.config.__dict__['sampling'] = config_temp.__dict__['sampling']	1✔
437	self.config.__dict__['neg_rate'] = config_temp.__dict__['neg_rate']	1✔
438
439	self.model.__init__(**self.config.__dict__)	1✔
440
441	self.model.load_state_dict(torch.load(str(model_path_file)))	1✔
442	self.model.eval()	1✔
443	else:
444	raise ValueError("Cannot load model from %s" % model_path_file)	×
445
446	def display(self):	1✔
447	"""Function to display embedding."""
448	options = {"ent_only_plot": True,	1✔
449	"rel_only_plot": not self.config.plot_entity_only,
450	"ent_and_rel_plot": not self.config.plot_entity_only}
451
452	if self.config.plot_embedding:	1!
453	viz = Visualization(self.model, self.config, vis_opts=options)	1✔
454	viz.plot_embedding(resultpath=self.config.path_figures, algos=self.model.model_name, show_label=False)	1✔
455
456	if self.config.plot_training_result:	1!
457	viz = Visualization(self.model, self.config)	1✔
458	viz.plot_train_result()	1✔
459
460	if self.config.plot_testing_result:	1!
461	viz = Visualization(self.model, self.config)	1✔
462	viz.plot_test_result()	1✔
463
464	def export_embeddings(self):	1✔
465	"""
466	Export embeddings in tsv and pandas pickled format.
467	With tsvs (both label, vector files), you can:
468	1) Use those pretained embeddings for your applications.
469	2) Visualize the embeddings in this website to gain insights. (https://projector.tensorflow.org/)
470
471	Pandas dataframes can be read with pd.read_pickle('desired_file.pickle')
472	"""
473	save_path = self.config.path_embeddings / self.model.model_name	1✔
474	save_path.mkdir(parents=True, exist_ok=True)	1✔
475
476	idx2ent = self.config.knowledge_graph.read_cache_data('idx2entity')	1✔
477	idx2rel = self.config.knowledge_graph.read_cache_data('idx2relation')	1✔
478
479	with open(str(save_path / "ent_labels.tsv"), 'w') as l_export_file:	1✔
480	for label in idx2ent.values():	1✔
481	l_export_file.write(label + "\n")	1✔
482
483	with open(str(save_path / "rel_labels.tsv"), 'w') as l_export_file:	1✔
484	for label in idx2rel.values():	1✔
485	l_export_file.write(label + "\n")	1✔
486
487	for named_embedding in self.model.parameter_list:	1✔
488	all_ids = list(range(0, int(named_embedding.weight.shape[0])))	1✔
489
490	stored_name = named_embedding.name	1✔
491
492	if len(named_embedding.weight.shape) == 2:	1!
493	all_embs = named_embedding.weight.detach().cpu().numpy()	1✔
494	with open(str(save_path / ("%s.tsv" % stored_name)), 'w') as v_export_file:	1✔
495	for idx in all_ids:	1✔
496	v_export_file.write("\t".join([str(x) for x in all_embs[idx]]) + "\n")	1✔
497
498	def save_training_result(self):	1✔
499	"""Function that saves training result"""
500	files = os.listdir(str(self.config.path_result))	1✔
501	l = len([f for f in files if self.model.model_name in f if 'Training' in f])	1✔
502	df = pd.DataFrame(self.training_results, columns=['Epochs', 'Loss'])	1✔
503	with open(str(self.config.path_result / (self.model.model_name + '_Training_results_' + str(l) + '.csv')),	1✔
504	'w') as fh:
505	df.to_csv(fh)	1✔

Sujit-O / pykg2vec / 759

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous