Implemented checkpointing in a different folder

Author: RandomDefaultUser

examples/advanced/ex01_checkpoint_training.py

@@ -32,6 +32,7 @@ def initial_setup():
     # as "ex07".
     parameters.running.checkpoints_each_epoch = 5
     parameters.running.checkpoint_name = "ex01_checkpoint"
+    parameters.running.checkpoint_path = "./"
 
     data_handler = mala.DataHandler(parameters)
     data_handler.add_snapshot(
@@ -62,9 +63,9 @@ def initial_setup():
     return parameters, test_network, data_handler, test_trainer
 
 
-if mala.Trainer.run_exists("ex01_checkpoint"):
+if mala.Trainer.run_exists("ex01_checkpoint", path="./"):
     parameters, network, datahandler, trainer = mala.Trainer.load_run(
-        "ex01_checkpoint"
+        "ex01_checkpoint", path="./"
     )
     printout("Starting resumed training.")
 else:

mala/common/parameters.py

@@ -881,6 +881,9 @@ class ParametersRunning(ParametersBase):
         Name used for the checkpoints. Using this, multiple runs
         can be performed in the same directory.
 
+    checkpoint_path : string
+        Path where the checkpoints will be saved (and loaded from)
+
     run_name : string
         Name of the run used for logging.
 
@@ -972,6 +975,7 @@ def __init__(self):
         self.checkpoints_each_epoch = 0
         # self.checkpoint_best_so_far = False
         self.checkpoint_name = "checkpoint_mala"
+        self.checkpoint_path = "./"
         self.run_name = ""
         self.logging_dir = "./mala_logging"
         self.logging_dir_append_date = True

mala/network/trainer.py

@@ -124,14 +124,23 @@ def __init__(self, params, network, data, _optimizer_dict=None):
         self._validation_graph = None
 
     @classmethod
-    def run_exists(cls, run_name, params_format="json", zip_run=True):
+    def run_exists(
+        cls, run_name, path="./", params_format="json", zip_run=True
+    ):
         """
         Check if a hyperparameter optimization checkpoint exists.
 
         Returns True if it does.
 
         Parameters
         ----------
+        path : str
+            Path to check for saved run.
+
+        zip_run : bool
+            If True, MALA will check for a .zip file. If False,
+            then separate files will be checked for.
+
         run_name : string
             Name of the checkpoint.
 
@@ -145,12 +154,14 @@ def run_exists(cls, run_name, params_format="json", zip_run=True):
 
         """
         if zip_run is True:
-            return os.path.isfile(run_name + ".zip")
+            return os.path.isfile(os.path.join(path, run_name + ".zip"))
         else:
-            network_name = run_name + ".network.pth"
-            iscaler_name = run_name + ".iscaler.pkl"
-            oscaler_name = run_name + ".oscaler.pkl"
-            param_name = run_name + ".params." + params_format
+            network_name = os.path.join(path, run_name + ".network.pth")
+            iscaler_name = os.path.join(path, run_name + ".iscaler.pkl")
+            oscaler_name = os.path.join(path, run_name + ".oscaler.pkl")
+            param_name = os.path.join(
+                path, run_name + ".params." + params_format
+            )
             optimizer_name = run_name + ".optimizer.pth"
             return all(
                 map(
@@ -1281,7 +1292,10 @@ def __create_training_checkpoint(self):
         Follows https://pytorch.org/tutorials/recipes/recipes/saving_and_
         loading_a_general_checkpoint.html to some degree.
         """
-        optimizer_name = self.parameters.checkpoint_name + ".optimizer.pth"
+        optimizer_name = os.path.join(
+            self.parameters.checkpoint_path,
+            self.parameters.checkpoint_name + ".optimizer.pth",
+        )
 
         # Next, we save all the other objects.
 
@@ -1306,14 +1320,11 @@ def __create_training_checkpoint(self):
         torch.save(
             save_dict, optimizer_name, _use_new_zipfile_serialization=False
         )
-        if self.parameters.run_name != "":
-            self.save_run(
-                self.parameters.checkpoint_name,
-                save_runner=True,
-                path=self.parameters.run_name,
-            )
-        else:
-            self.save_run(self.parameters.checkpoint_name, save_runner=True)
+        self.save_run(
+            self.parameters.checkpoint_name,
+            save_runner=True,
+            path=self.parameters.checkpoint_path,
+        )
 
     @staticmethod
     def __average_validation(val, device="cpu"):