create packaged models from checkpoint

README.md

@@ -328,7 +328,7 @@ import torch
 import kge.model
 
 # download link for this checkpoint given under results above
-model = kge.model.KgeModel.load_from_checkpoint('fb15k-237-rescal.pt')
+model = kge.model.KgeModel.load_from('fb15k-237-rescal.pt')
 
 s = torch.Tensor([0, 2,]).long()             # subject indexes
 p = torch.Tensor([0, 1,]).long()             # relation indexes

kge/cli.py

@@ -10,6 +10,7 @@
 from kge.job import Job
 from kge.misc import get_git_revision_short_hash, kge_base_dir, is_number
 from kge.util.dump import add_dump_parsers, dump
+from kge.util.package import package_model, add_package_parser
 
 
 def argparse_bool_type(v):
@@ -130,6 +131,7 @@ def create_parser(config, additional_args=[]):
             default="default",
         )
     add_dump_parsers(subparsers)
+    add_package_parser(subparsers)
     return parser
 
 
@@ -163,6 +165,11 @@ def main():
         dump(args)
         exit()
 
+    # package command
+    if args.command == "package":
+        package_model(args.checkpoint)
+        exit()
+
     # start command
     if args.command == "start":
         # use toy config file if no config given
@@ -272,9 +279,10 @@ def main():
             dataset = Dataset.load(config)
 
             # let's go
-            job = Job.create(config, dataset)
             if args.command == "resume":
-                job.resume(checkpoint_file)
+                job = Job.resume(checkpoint_file, config=config)
+            else:
+                job = Job.create(config, dataset)
             job.run()
     except BaseException as e:
         tb = traceback.format_exc()

kge/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import collections
 import copy
 import datetime
@@ -423,6 +425,26 @@ def init_folder(self):
             return True
         return False
 
+    @staticmethod
+    def load_from(checkpoint_config: Config, new_config=None, folder=None) -> Config:
+        """
+        Load a config from checkpoint and update all deprecated keys.
+        Overwrite options in checkpoint config with a new config.
+        Args:
+            checkpoint_config: Config stored in checkpoint
+            new_config: new config with options to overwrite
+
+        Returns: Config object
+
+        """
+        config = Config()  # round trip to handle deprecated configs
+        config.load_options(checkpoint_config.options)
+        if folder is not None:
+            config.folder = folder
+        if new_config is not None:
+            config.load_options(new_config.options)
+        return config
+
     def checkpoint_file(self, cpt_id: Union[str, int]) -> str:
         "Return path of checkpoint file for given checkpoint id"
         from kge.misc import is_number

kge/dataset.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import csv
 import os
 import sys
@@ -91,6 +93,20 @@ def load(config: Config, preload_data=True):
                 dataset.split(split)
         return dataset
 
+    @staticmethod
+    def load_meta(config: Config, meta_checkpoint: Dict, preload_data=True) -> Dataset:
+        """Loads a dataset and overrides meta data."""
+        dataset = Dataset.load(config, preload_data)
+        dataset._meta.update(meta_checkpoint)
+        return dataset
+
+    def save_meta(self, keys: List[str]):
+        """Creates a dataset_meta dictionary for a checkpoint."""
+        meta_checkpoint = {}
+        for key in keys:
+            meta_checkpoint[key] = self._meta[key]
+        return meta_checkpoint
+
     @staticmethod
     def _to_valid_filename(s):
         invalid_chars = "\n\t\\/"

kge/job/auto_search.py

@@ -49,24 +49,16 @@ def save(self, filename):
             filename,
         )
 
-    def resume(self, checkpoint_file=None):
-        if checkpoint_file is None:
-            last_checkpoint = self.config.last_checkpoint()
-            if last_checkpoint is not None:
-                checkpoint_file = self.config.checkpoint_file(last_checkpoint)
-
-        if checkpoint_file is not None:
-            self.resumed_from_job_id = self.load(checkpoint_file)
-            self.trace(
-                event="job_resumed", checkpoint_file=checkpoint_file
-            )
-            self.config.log(
-                "Resumed from {} of job {}".format(
-                    checkpoint_file, self.resumed_from_job_id
-                )
+    def load(self, checkpoint):
+        self.resumed_from_job_id = checkpoint.get("job_id")
+        self.trace(
+            event="job_resumed", checkpoint_file=checkpoint["file"]
+        )
+        self.config.log(
+            "Resumed from {} of job {}".format(
+                checkpoint["file"], self.resumed_from_job_id
             )
-        else:
-            self.config.log("No checkpoint found, starting from scratch...")
+        )
 
     # -- Abstract methods --------------------------------------------------------------
 

kge/job/eval.py

@@ -1,6 +1,10 @@
 import torch
 
+from kge import Config, Dataset
 from kge.job import Job
+from kge.model import KgeModel
+
+from typing import Dict, Union
 
 
 class EvaluationJob(Job):
@@ -83,20 +87,27 @@ def run(self) -> dict:
         """ Compute evaluation metrics, output results to trace file """
         raise NotImplementedError
 
-    def resume(self, checkpoint_file=None):
-        """Load model state from last or specified checkpoint."""
-        # load model
-        from kge.job import TrainingJob
-
-        training_job = TrainingJob.create(self.config, self.dataset)
-        training_job.resume(checkpoint_file)
-        self.model = training_job.model
-        self.epoch = training_job.epoch
-        self.resumed_from_job_id = training_job.resumed_from_job_id
+    def load(self, checkpoint, model):
+        self.resumed_from_job_id = checkpoint.get("job_id")
+        self.epoch = checkpoint["epoch"]
         self.trace(
-            event="job_resumed", epoch=self.epoch, checkpoint_file=checkpoint_file
+            event="job_resumed", epoch=self.epoch, checkpoint_file=checkpoint["file"]
         )
 
+    @classmethod
+    def load_from(
+        cls,
+        checkpoint: Union[str, Dict],
+        config: Config = None,
+        dataset: Dataset = None,
+        parent_job=None,
+    ) -> Job:
+        if config is None:
+            config = Config()
+            config.set("job.type", "eval")
+        eval_job = super().load_from(checkpoint, config, dataset, parent_job=parent_job)
+        return eval_job
+
 
 # HISTOGRAM COMPUTATION ###############################################################
 

kge/job/grid_search.py

@@ -69,7 +69,6 @@ def run(self):
         # and run it
         if self.config.get("grid_search.run"):
             job = Job.create(self.config, self.dataset, parent_job=self)
-            job.resume()
             job.run()
         else:
             self.config.log("Skipping running of search job as requested by user...")

kge/job/job.py

@@ -1,10 +1,14 @@
+from __future__ import annotations
+
 from kge import Config, Dataset
+from kge.model import KgeModel
+from kge.util import load_checkpoint
 import uuid
 
 from kge.misc import get_git_revision_short_hash
 import os
 import socket
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Union
 
 
 def _trace_job_creation(job: "Job"):
@@ -55,36 +59,84 @@ def __init__(self, config: Config, dataset: Dataset, parent_job: "Job" = None):
             for f in Job.job_created_hooks:
                 f(self)
 
-    def resume(self, checkpoint_file: str = None):
-        """Load job state from last or specified checkpoint.
-
-        Restores all relevant state to resume a previous job. To run the restored job,
-        use :func:`run`.
-
-        Should set `resumed_from_job` to the job ID of the previous job.
-
-        """
-        raise NotImplementedError
-
-    def run(self):
-        raise NotImplementedError
-
-    def create(config: Config, dataset: Dataset, parent_job: "Job" = None) -> "Job":
-        """Creates a job for a given configuration."""
-
+    @staticmethod
+    def create(config: Config, dataset: Dataset, parent_job=None, model=None):
         from kge.job import TrainingJob, EvaluationJob, SearchJob
 
-        if config.get("job.type") == "train":
-            job = TrainingJob.create(config, dataset, parent_job)
-        elif config.get("job.type") == "search":
-            job = SearchJob.create(config, dataset, parent_job)
-        elif config.get("job.type") == "eval":
-            job = EvaluationJob.create(config, dataset, parent_job)
+        job_type = config.get("job.type")
+        if job_type == "train":
+            return TrainingJob.create(
+                config, dataset, parent_job=parent_job, model=model
+            )
+        elif job_type == "search":
+            return SearchJob.create(config, dataset, parent_job=parent_job)
+        elif job_type == "eval":
+            return EvaluationJob.create(
+                config, dataset, parent_job=parent_job, model=model
+            )
         else:
             raise ValueError("unknown job type")
 
+    @staticmethod
+    def resume(
+        checkpoint: str = None,
+        config: Config = None,
+        dataset: Dataset = None,
+        parent_job=None,
+    ) -> Job:
+        if checkpoint is None and config is None:
+            raise ValueError(
+                "Please provide either the config file located in the folder structure "
+                "containing the checkpoint or the checkpoint itself."
+            )
+        elif checkpoint is None:
+            last_checkpoint = config.last_checkpoint()
+            if last_checkpoint is not None:
+                checkpoint = config.checkpoint_file(last_checkpoint)
+
+        if checkpoint is not None:
+            job = Job.load_from(checkpoint, config, dataset, parent_job=parent_job)
+            if type(checkpoint) == str:
+                job.config.log("Loading checkpoint from {}...".format(checkpoint))
+        else:
+            job = Job.create(config, dataset, parent_job=parent_job)
+            job.config.log("No checkpoint found, starting from scratch...")
+        return job
+
+    @classmethod
+    def load_from(
+        cls,
+        checkpoint: Union[str, Dict],
+        config: Config = None,
+        dataset: Dataset = None,
+        parent_job=None,
+    ) -> Job:
+        if config is not None:
+            device = config.get("job.device")
+        else:
+            device = "cpu"
+        if type(checkpoint) == str:
+            checkpoint = load_checkpoint(checkpoint, device)
+        config = Config.load_from(
+            checkpoint["config"], config, folder=checkpoint["folder"]
+        )
+        model: KgeModel = None
+        dataset = None
+        if checkpoint["model"] is not None:
+            model = KgeModel.load_from(checkpoint, config=config, dataset=dataset)
+            dataset = model.dataset
+        if dataset is None:
+            dataset = Dataset.load(config)
+        job = Job.create(config, dataset, parent_job, model)
+        job.load(checkpoint, model)
         return job
 
+    def load(self, checkpoint, model):
+        pass
+
+    def run(self):
+        raise NotImplementedError
+
     def trace(self, **kwargs) -> Dict[str, Any]:
         """Write a set of key-value pairs to the trace file and automatically append
         information about this job. See `Config.trace` for more information."""

kge/job/manual_search.py

@@ -32,10 +32,6 @@ def __init__(self, config: Config, dataset: Dataset, parent_job=None):
             for f in Job.job_created_hooks:
                 f(self)
 
-    def resume(self, checkpoint_file=None):
-        # no need to do anything here; run code automatically resumes
-        pass
-
     def run(self):
         # read search configurations and expand them to full configs
         search_configs = copy.deepcopy(self.config.get("manual_search.configurations"))

kge/job/search.py

@@ -41,6 +41,7 @@ def __init__(self, config, dataset, parent_job=None):
             for f in Job.job_created_hooks:
                 f(self)
 
+    @staticmethod
     def create(config, dataset, parent_job=None):
         """Factory method to create a search job."""
 
@@ -130,8 +131,12 @@ def _run_train_job(sicnk, device=None):
                 train_job_config.get("job.device"),
             )
         )
-        job = Job.create(train_job_config, search_job.dataset, parent_job=search_job)
-        job.resume()
+        job = Job.resume(
+            checkpoint=None,
+            config=train_job_config,
+            dataset=search_job.dataset,
+            parent_job=search_job,
+        )
 
         # process the trace entries to far (in case of a resumed job)
         metric_name = search_job.config.get("valid.metric")