Add loss/cost on validation data https://github.com/uma-pi1/kge/issues/2

kge/config-default.yaml

@@ -235,6 +235,10 @@ KvsAll:
     s_o: False
     _po: True
 
+  # Dataset splits from which the labels for a query are taken from. Default: If
+  # nothing is specified, then the train split is used.
+  label_splits: []
+
 # Options for negative sampling training (train.type=="negative_sampling")
 negative_sampling:
   # Negative sampler to use
@@ -264,7 +268,7 @@ negative_sampling:
     p: False       # as above
     o: False       # as above
 
-    split: ''      # split containing the positives; default is train.split
+    splits: []     # splits containing the positives; default is train.split
 
     # Implementation to use for filtering.
     # standard: use slow generic implementation, available for all samplers
@@ -327,7 +331,8 @@ eval:
   # mean_reciprocal_rank_filtered_with_test.
   filter_with_test: True
 
-  # Type of evaluation (entity_ranking only at the moment)
+  # Type of evaluation (entity_ranking or training_loss). Currently, 
+  # entity_ranking runs training_loss as well.
   type: entity_ranking
 
   # Compute Hits@K for these choices of K

kge/indexing.py

@@ -2,6 +2,7 @@
 from collections import defaultdict, OrderedDict
 import numba
 import numpy as np
+from kge.misc import powerset, merge_dicts_of_1dim_torch_tensors
 
 
 def _group_by(keys, values) -> dict:
@@ -222,13 +223,31 @@ def _invert_ids(dataset, obj: str):
     dataset.config.log(f"Indexed {len(inv)} {obj} ids", prefix="  ")
 
 
+def merge_KvsAll_indexes(dataset, split, key):
+    value = dict([("sp", "o"), ("po", "s"), ("so", "p")])[key]
+    split_combi_str = "_".join(sorted(split))
+    index_name = f"{split_combi_str}_{key}_to_{value}"
+    indexes = [dataset.index(f"{_split}_{key}_to_{value}") for _split in split]
+    dataset._indexes[index_name] = merge_dicts_of_1dim_torch_tensors(indexes)
+    return dataset._indexes[index_name]
+
+
 def create_default_index_functions(dataset: "Dataset"):
     for split in dataset.files_of_type("triples"):
         for key, value in [("sp", "o"), ("po", "s"), ("so", "p")]:
             # self assignment needed to capture the loop var
             dataset.index_functions[f"{split}_{key}_to_{value}"] = IndexWrapper(
                 index_KvsAll, split=split, key=key
             )
+    # create all combinations of splits of length 2 and 3
+    for split_combi in powerset(dataset.files_of_type("triples"), [2, 3]):
+        for key, value in [("sp", "o"), ("po", "s"), ("so", "p")]:
+            split_combi_str = "_".join(sorted(split_combi))
+            index_name = f"{split_combi_str}_{key}_to_{value}"
+            dataset.index_functions[index_name] = IndexWrapper(
+                merge_KvsAll_indexes, split=split_combi, key=key
+            )
+
     dataset.index_functions["relation_types"] = index_relation_types
     dataset.index_functions["relations_per_type"] = index_relation_types
     dataset.index_functions["frequency_percentiles"] = index_frequency_percentiles

kge/job/entity_ranking.py

@@ -1,9 +1,10 @@
 import math
 import time
+from typing import Dict, Any
 
 import torch
 import kge.job
-from kge.job import EvaluationJob, Job
+from kge.job import EvaluationJob, Job, TrainingJob
 from kge import Config, Dataset
 from collections import defaultdict
 
@@ -13,18 +14,23 @@ class EntityRankingJob(EvaluationJob):
 
     def __init__(self, config: Config, dataset: Dataset, parent_job, model):
         super().__init__(config, dataset, parent_job, model)
-        self.is_prepared = False
 
         if self.__class__ == EntityRankingJob:
             for f in Job.job_created_hooks:
                 f(self)
 
+        max_k = min(
+            self.dataset.num_entities(), max(self.config.get("eval.hits_at_k_s"))
+        )
+        self.hits_at_k_s = list(
+            filter(lambda x: x <= max_k, self.config.get("eval.hits_at_k_s"))
+        )
+        self.filter_with_test = config.get("eval.filter_with_test")
+
+
     def _prepare(self):
         """Construct all indexes needed to run."""
 
-        if self.is_prepared:
-            return
-
         # create data and precompute indexes
         self.triples = self.dataset.split(self.config.get("eval.split"))
         for split in self.filter_splits:
@@ -75,16 +81,8 @@ def _collate(self, batch):
         return batch, label_coords, test_label_coords
 
     @torch.no_grad()
-    def run(self) -> dict:
-        self._prepare()
-
-        was_training = self.model.training
-        self.model.eval()
-        self.config.log(
-            "Evaluating on "
-            + self.eval_split
-            + " data (epoch {})...".format(self.epoch)
-        )
+    def _run(self) -> Dict[str, Any]:
+
         num_entities = self.dataset.num_entities()
 
         # we also filter with test data if requested
@@ -399,28 +397,6 @@ def merge_hist(target_hists, source_hists):
             event="eval_completed",
             **metrics,
         )
-        for f in self.post_epoch_trace_hooks:
-            f(self, trace_entry)
-
-        # if validation metric is not present, try to compute it
-        metric_name = self.config.get("valid.metric")
-        if metric_name not in trace_entry:
-            trace_entry[metric_name] = eval(
-                self.config.get("valid.metric_expr"),
-                None,
-                dict(config=self.config, **trace_entry),
-            )
-
-        # write out trace
-        trace_entry = self.trace(**trace_entry, echo=True, echo_prefix="  ", log=True)
-
-        # reset model and return metrics
-        if was_training:
-            self.model.train()
-        self.config.log("Finished evaluating on " + self.eval_split + " split.")
-
-        for f in self.post_valid_hooks:
-            f(self, trace_entry)
 
         return trace_entry
 

kge/job/eval.py

@@ -1,10 +1,10 @@
-import torch
+import time
+from typing import Any, Optional, Dict
 
+import torch
 from kge import Config, Dataset
-from kge.job import Job
-from kge.model import KgeModel
 
-from typing import Dict, Union, Optional
+from kge.job import Job, TrainingJob
 
 
 class EvaluationJob(Job):
@@ -16,12 +16,6 @@ def __init__(self, config, dataset, parent_job, model):
         self.model = model
         self.batch_size = config.get("eval.batch_size")
         self.device = self.config.get("job.device")
-        max_k = min(
-            self.dataset.num_entities(), max(self.config.get("eval.hits_at_k_s"))
-        )
-        self.hits_at_k_s = list(
-            filter(lambda x: x <= max_k, self.config.get("eval.hits_at_k_s"))
-        )
         self.config.check("train.trace_level", ["example", "batch", "epoch"])
         self.trace_examples = self.config.get("eval.trace_level") == "example"
         self.trace_batch = (
@@ -31,9 +25,11 @@ def __init__(self, config, dataset, parent_job, model):
         self.filter_splits = self.config.get("eval.filter_splits")
         if self.eval_split not in self.filter_splits:
             self.filter_splits.append(self.eval_split)
-        self.filter_with_test = config.get("eval.filter_with_test")
         self.epoch = -1
 
+        self.verbose = True
+        self.is_prepared = False
+
         #: Hooks run after training for an epoch.
         #: Signature: job, trace_entry
         self.post_epoch_hooks = []
@@ -64,6 +60,22 @@ def __init__(self, config, dataset, parent_job, model):
         if config.get("eval.metrics_per.argument_frequency"):
             self.hist_hooks.append(hist_per_frequency_percentile)
 
+        # Add the training loss as a default to every evaluation job
+        # TODO: create AggregatingEvaluationsJob that runs and aggregates a list
+        #  of EvaluationAjobs, such that users can configure combinations of
+        #  EvalJobs themselves. Then this can be removed.
+        #  See https://github.com/uma-pi1/kge/issues/102
+        if not isinstance(self, TrainingLossEvaluationJob):
+            self.eval_train_loss_job = TrainingLossEvaluationJob(
+                config, dataset, parent_job=self, model=model
+            )
+            self.eval_train_loss_job.verbose = False
+            self.post_epoch_trace_hooks.append(
+                lambda job, trace: trace.update(
+                    avg_loss=self.eval_train_loss_job.run()["avg_loss"]
+                )
+            )
+
         # all done, run job_created_hooks if necessary
         if self.__class__ == EvaluationJob:
             for f in Job.job_created_hooks:
@@ -81,10 +93,66 @@ def create(config, dataset, parent_job=None, model=None):
             return EntityPairRankingJob(
                 config, dataset, parent_job=parent_job, model=model
             )
+        elif config.get("eval.type") == "training_loss":
+            return TrainingLossEvaluationJob(
+                config, dataset, parent_job=parent_job, model=model
+            )
         else:
             raise ValueError("eval.type")
 
-    def run(self) -> dict:
+    def _prepare(self):
+        """Prepare this job for running. Guaranteed to be called exactly once
+        """
+        raise NotImplementedError
+
+    def run(self) -> Dict[str, Any]:
+
+        if not self.is_prepared:
+            self._prepare()
+            self.model.prepare_job(self)  # let the model add some hooks
+            self.is_prepared = True
+
+        was_training = self.model.training
+        self.model.eval()
+        self.config.log(
+            "Evaluating on "
+            + self.eval_split
+            + " data (epoch {})...".format(self.epoch),
+            echo=self.verbose,
+        )
+
+        trace_entry = self._run()
+
+        # if validation metric is not present, try to compute it
+        metric_name = self.config.get("valid.metric")
+        if metric_name not in trace_entry:
+            trace_entry[metric_name] = eval(
+                self.config.get("valid.metric_expr"),
+                None,
+                dict(config=self.config, **trace_entry),
+            )
+
+        for f in self.post_epoch_trace_hooks:
+            f(self, trace_entry)
+
+        # write out trace
+        trace_entry = self.trace(
+            **trace_entry, echo=self.verbose, echo_prefix="  ", log=True
+        )
+
+        # reset model and return metrics
+        if was_training:
+            self.model.train()
+        self.config.log(
+            "Finished evaluating on " + self.eval_split + " split.", echo=self.verbose
+        )
+
+        for f in self.post_valid_hooks:
+            f(self, trace_entry)
+
+        return trace_entry
+
+    def _run(self) -> Dict[str, Any]:
         """ Compute evaluation metrics, output results to trace file """
         raise NotImplementedError
 
@@ -132,8 +200,65 @@ def create_from(
         return super().create_from(checkpoint, new_config, dataset, parent_job)
 
 
-# HISTOGRAM COMPUTATION ###############################################################
+class TrainingLossEvaluationJob(EvaluationJob):
+    """ Entity ranking evaluation protocol """
+
+    def __init__(self, config: Config, dataset: Dataset, parent_job, model):
+        super().__init__(config, dataset, parent_job, model)
+        self.is_prepared = True
+
+        train_job_on_eval_split_config = config.clone()
+        train_job_on_eval_split_config.set("train.split", self.eval_split)
+        train_job_on_eval_split_config.set("verbose", False)
+        train_job_on_eval_split_config.set(
+            "negative_sampling.filtering.splits",
+            [self.config.get("train.split"), self.eval_split] + ["valid"]
+            if self.eval_split == "test"
+            else [],
+        )
+        train_job_on_eval_split_config.set(
+            "KvsAll.label_splits",
+            [self.config.get("train.split"), self.eval_split] + ["valid"]
+            if self.eval_split == "test"
+            else [],
+        )
+        self._train_job = TrainingJob.create(
+            config=train_job_on_eval_split_config,
+            parent_job=self,
+            dataset=dataset,
+            model=model,
+            initialize_for_forward_only=True,
+        )
+        self._train_job_verbose = False
 
+        if self.__class__ == TrainingLossEvaluationJob:
+            for f in Job.job_created_hooks:
+                f(self)
+
+    @torch.no_grad()
+    def _run(self) -> Dict[str, Any]:
+
+        epoch_time = -time.time()
+
+        self.epoch = self.parent_job.epoch
+        epoch_time += time.time()
+
+        train_trace_entry = self._train_job.run_epoch()
+        # compute trace
+        trace_entry = dict(
+            type="training_loss",
+            scope="epoch",
+            split=self.eval_split,
+            epoch=self.epoch,
+            epoch_time=epoch_time,
+            event="eval_completed",
+            avg_loss=train_trace_entry["avg_loss"],
+        )
+
+        return trace_entry
+
+
+# HISTOGRAM COMPUTATION ###############################################################
 
 def __initialize_hist(hists, key, job):
     """If there is no histogram with given `key` in `hists`, add an empty one."""