Log best epoch when loading checkpoint (#699)

---------

Co-authored-by: Filippo Bigi <[email protected]>

Author: PicoCentauri

src/metatrain/cli/eval.py

@@ -238,7 +238,7 @@ def _eval_targets(
     mean_per_atom = np.mean(timings_per_atom)
     std_per_atom = np.std(timings_per_atom)
     logging.info(
-        f"evaluation time: {total_time:.2f} s "
+        f"Evaluation time: {total_time:.2f} s "
         f"[{1000.0 * mean_per_atom:.4f} ± "
         f"{1000.0 * std_per_atom:.4f} ms per atom]"
     )

src/metatrain/cli/train.py

@@ -568,7 +568,11 @@ def train_model(
         trainer.save_checkpoint(model, checkpoint_output)
     except Exception as e:
         raise ArchitectureError(e)
+
     if checkpoint_output.exists():
+        # Reload ensuring (best) model intended for inference
+        model = load_model(checkpoint_output)
+
         logging.info(f"Final checkpoint: {checkpoint_output.absolute().resolve()}")
 
     mts_atomistic_model = model.export()

src/metatrain/pet/checkpoints.py

@@ -1,25 +1,42 @@
-def model_update_v1_v2(state_dict):
-    # This if-statement is necessary to handle cases when
-    # best_model_state_dict and model_state_dict are the same.
-    # In that case, the both are updated within the first call of
-    # this function in the PET.update_checkpoint() method.
-    if (
-        state_dict is not None
-        and "additive_models.0.model.type_to_index" not in state_dict
-    ):
-        state_dict["additive_models.0.model.type_to_index"] = state_dict.pop(
-            "additive_models.0.type_to_index"
-        )
+###########################
+# MODEL ###################
+###########################
+
+
+def model_update_v1_v2(checkpoint):
+    for key in ["model_state_dict", "best_model_state_dict"]:
+        if (state_dict := checkpoint.get(key)) is not None:
+            state_dict["additive_models.0.model.type_to_index"] = state_dict.pop(
+                "additive_models.0.type_to_index"
+            )
+
+
+def model_update_v2_v3(checkpoint):
+    for key in ["model_state_dict", "best_model_state_dict"]:
+        if (state_dict := checkpoint.get(key)) is not None:
+            if "train_hypers" in state_dict:
+                finetune_config = state_dict["train_hypers"].get("finetune", {})
+            else:
+                finetune_config = {}
+            state_dict["finetune_config"] = finetune_config
+
+
+def model_update_v3_v4(checkpoint):
+    checkpoint["epoch"] = checkpoint.get("epoch")
+    checkpoint["best_epoch"] = checkpoint.get("best_epoch")
+
+    if checkpoint["best_model_state_dict"] is not None:
+        checkpoint["best_model_state_dict"] = checkpoint.get("best_model_state_dict")
+
+
+###########################
+# TRAINER #################
+###########################
 
 
 def trainer_update_v1_v2(checkpoint):
     checkpoint["train_hypers"] = checkpoint["train_hypers"].get("scheduler_factor", 0.5)
 
 
-def model_update_v2_v3(state_dict):
-    if state_dict is not None:
-        if "train_hypers" in state_dict:
-            finetune_config = state_dict["train_hypers"].get("finetune", {})
-        else:
-            finetune_config = {}
-        state_dict["finetune_config"] = finetune_config
+def trainer_update_v2_v3(checkpoint):
+    checkpoint["best_epoch"] = checkpoint.get("best_epoch")

src/metatrain/pet/model.py

@@ -1,3 +1,4 @@
+import logging
 import warnings
 from math import prod
 from typing import Any, Dict, List, Literal, Optional
@@ -40,7 +41,7 @@ class PET(ModelInterface):
 
     """
 
-    __checkpoint_version__ = 3
+    __checkpoint_version__ = 4
     __supported_devices__ = ["cuda", "cpu"]
     __supported_dtypes__ = [torch.float32, torch.float64]
     __default_metadata__ = ModelMetadata(
@@ -686,25 +687,23 @@ def load_checkpoint(
         checkpoint: Dict[str, Any],
         context: Literal["restart", "finetune", "export"],
     ) -> "PET":
-        model_data = checkpoint["model_data"]
-
         if context == "restart":
+            logging.info(f"Using latest model from epoch {checkpoint['epoch']}")
             model_state_dict = checkpoint["model_state_dict"]
-        elif context == "finetune" or context == "export":
+        elif context in {"finetune", "export"}:
+            logging.info(f"Using best model from epoch {checkpoint['best_epoch']}")
             model_state_dict = checkpoint["best_model_state_dict"]
-            if model_state_dict is None:
-                model_state_dict = checkpoint["model_state_dict"]
         else:
             raise ValueError("Unknown context tag for checkpoint loading!")
 
-        finetune_config = model_state_dict.pop("finetune_config", {})
-
         # Create the model
+        model_data = checkpoint["model_data"]
         model = cls(
             hypers=model_data["model_hypers"],
             dataset_info=model_data["dataset_info"],
         )
 
+        finetune_config = model_state_dict.pop("finetune_config", {})
         if finetune_config:
             # Apply the finetuning strategy
             model = apply_finetuning_strategy(model, finetune_config)
@@ -890,14 +889,11 @@ def _get_system_indices_and_labels(
 
     @classmethod
     def upgrade_checkpoint(cls, checkpoint: Dict) -> Dict:
-        if checkpoint["model_ckpt_version"] == 1:
-            checkpoints.model_update_v1_v2(checkpoint["model_state_dict"])
-            checkpoints.model_update_v1_v2(checkpoint["best_model_state_dict"])
-            checkpoint["model_ckpt_version"] = 2
-        if checkpoint["model_ckpt_version"] == 2:
-            checkpoints.model_update_v2_v3(checkpoint["model_state_dict"])
-            checkpoints.model_update_v2_v3(checkpoint["best_model_state_dict"])
-            checkpoint["model_ckpt_version"] = 3
+        for v in range(1, cls.__checkpoint_version__):
+            if checkpoint["model_ckpt_version"] == v:
+                update = getattr(checkpoints, f"model_update_v{v}_v{v + 1}")
+                update(checkpoint)
+                checkpoint["model_ckpt_version"] = v + 1
 
         if checkpoint["model_ckpt_version"] != cls.__checkpoint_version__:
             raise RuntimeError(
@@ -919,7 +915,9 @@ def get_checkpoint(self) -> Dict:
                 "model_hypers": self.hypers,
                 "dataset_info": self.dataset_info,
             },
+            "epoch": None,
+            "best_epoch": None,
             "model_state_dict": model_state_dict,
-            "best_model_state_dict": None,
+            "best_model_state_dict": self.state_dict(),
         }
         return checkpoint

src/metatrain/pet/tests/checkpoints/model-v4_trainer-v2.ckpt.gz

@@ -1,4 +1,5 @@
 import copy
+import logging
 
 import pytest
 import torch
@@ -81,7 +82,7 @@ def model_trainer():
 
 
 @pytest.mark.parametrize("context", ["finetune", "restart", "export"])
-def test_get_checkpoint(context):
+def test_get_checkpoint(context, caplog):
     """
     Test that the checkpoint created by the model.get_checkpoint()
     function can be loaded back in all possible contexts.
@@ -93,8 +94,15 @@ def test_get_checkpoint(context):
     )
     model = PET(MODEL_HYPERS, dataset_info)
     checkpoint = model.get_checkpoint()
+
+    caplog.set_level(logging.INFO)
     PET.load_checkpoint(checkpoint, context)
 
+    if context == "restart":
+        assert "Using latest model from epoch None" in caplog.text
+    else:
+        assert "Using best model from epoch None" in caplog.text
+
 
 @pytest.mark.parametrize("cls_type", ["model", "trainer"])
 def test_failed_checkpoint_upgrade(cls_type):

src/metatrain/pet/tests/checkpoints/model-v4_trainer-v3.ckpt.gz

@@ -52,14 +52,15 @@ def func_lr_scheduler(epoch):
 
 
 class Trainer(TrainerInterface):
-    __checkpoint_version__ = 2
+    __checkpoint_version__ = 3
 
     def __init__(self, hypers):
         super().__init__(hypers)
 
         self.optimizer_state_dict = None
         self.scheduler_state_dict = None
         self.epoch = None
+        self.best_epoch = None
         self.best_metric = None
         self.best_model_state_dict = None
         self.best_optimizer_state_dict = None
@@ -520,6 +521,7 @@ def train(
                 self.best_model_state_dict = copy.deepcopy(
                     (model.module if is_distributed else model).state_dict()
                 )
+                self.best_epoch = epoch
                 self.best_optimizer_state_dict = copy.deepcopy(optimizer.state_dict())
 
             if epoch % self.hypers["checkpoint_interval"] == 0:
@@ -553,6 +555,7 @@ def save_checkpoint(self, model, path: Union[str, Path]):
                 "epoch": self.epoch,
                 "optimizer_state_dict": self.optimizer_state_dict,
                 "scheduler_state_dict": self.scheduler_state_dict,
+                "best_epoch": self.best_epoch,
                 "best_metric": self.best_metric,
                 "best_model_state_dict": self.best_model_state_dict,
                 "best_optimizer_state_dict": self.best_optimizer_state_dict,
@@ -570,29 +573,25 @@ def load_checkpoint(
         hypers: Dict[str, Any],
         context: Literal["restart", "finetune"],
     ) -> "Trainer":
-        epoch = checkpoint["epoch"]
-        optimizer_state_dict = checkpoint["optimizer_state_dict"]
-        scheduler_state_dict = checkpoint["scheduler_state_dict"]
-        best_metric = checkpoint["best_metric"]
-        best_model_state_dict = checkpoint["best_model_state_dict"]
-        best_optimizer_state_dict = checkpoint["best_optimizer_state_dict"]
-
-        # Create the trainer
         trainer = cls(hypers)
-        trainer.optimizer_state_dict = optimizer_state_dict
-        trainer.scheduler_state_dict = scheduler_state_dict
-        trainer.epoch = epoch
-        trainer.best_metric = best_metric
-        trainer.best_model_state_dict = best_model_state_dict
-        trainer.best_optimizer_state_dict = best_optimizer_state_dict
+        trainer.optimizer_state_dict = checkpoint["optimizer_state_dict"]
+        trainer.scheduler_state_dict = checkpoint["scheduler_state_dict"]
+        trainer.epoch = checkpoint["epoch"]
+        trainer.best_epoch = checkpoint["best_epoch"]
+        trainer.best_metric = checkpoint["best_metric"]
+        trainer.best_model_state_dict = checkpoint["best_model_state_dict"]
+        trainer.best_optimizer_state_dict = checkpoint["best_optimizer_state_dict"]
 
         return trainer
 
     @classmethod
     def upgrade_checkpoint(cls, checkpoint: Dict) -> Dict:
-        if checkpoint["trainer_ckpt_version"] == 1:
-            checkpoints.trainer_update_v1_v2(checkpoint)
-            checkpoint["trainer_ckpt_version"] = 2
+        for v in range(1, cls.__checkpoint_version__):
+            if checkpoint["trainer_ckpt_version"] == v:
+                update = getattr(checkpoints, f"trainer_update_v{v}_v{v + 1}")
+                update(checkpoint)
+                checkpoint["trainer_ckpt_version"] = v + 1
+
         if checkpoint["trainer_ckpt_version"] != cls.__checkpoint_version__:
             raise RuntimeError(
                 f"Unable to upgrade the checkpoint: the checkpoint is using "

src/metatrain/pet/tests/test_checkpoints.py

@@ -1,12 +1,28 @@
-def update_v1_v2(state_dict):
-    # This if-statement is necessary to handle cases when
-    # best_model_state_dict and model_state_dict are the same.
-    # In that case, the both are updated within the first call of
-    # this function in the PET.update_checkpoint() method.
-    if (
-        state_dict is not None
-        and "additive_models.0.model.type_to_index" not in state_dict
-    ):
-        state_dict["additive_models.0.model.type_to_index"] = state_dict.pop(
-            "additive_models.0.type_to_index"
-        )
+###########################
+# MODEL ###################
+###########################
+
+
+def model_update_v1_v2(checkpoint):
+    for key in ["model_state_dict", "best_model_state_dict"]:
+        if (state_dict := checkpoint.get(key)) is not None:
+            state_dict["additive_models.0.model.type_to_index"] = state_dict.pop(
+                "additive_models.0.type_to_index"
+            )
+
+
+def model_update_v2_v3(checkpoint):
+    checkpoint["epoch"] = checkpoint.get("epoch")
+    checkpoint["best_epoch"] = checkpoint.get("best_epoch")
+
+    if checkpoint["best_model_state_dict"] is not None:
+        checkpoint["best_model_state_dict"] = checkpoint.get("best_model_state_dict")
+
+
+###########################
+# TRAINER #################
+###########################
+
+
+def trainer_update_v1_v2(checkpoint):
+    checkpoint["best_epoch"] = checkpoint.get("best_epoch")

src/metatrain/pet/trainer.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Any, Dict, List, Literal, Optional
 
 import metatensor.torch as mts
@@ -171,7 +172,7 @@ def concatenate_structures(
 
 
 class SoapBpnn(ModelInterface):
-    __checkpoint_version__ = 2
+    __checkpoint_version__ = 3
     __supported_devices__ = ["cuda", "cpu"]
     __supported_dtypes__ = [torch.float64, torch.float32]
     __default_metadata__ = ModelMetadata(
@@ -670,18 +671,17 @@ def load_checkpoint(
         checkpoint: Dict[str, Any],
         context: Literal["restart", "finetune", "export"],
     ) -> "SoapBpnn":
-        model_data = checkpoint["model_data"]
-
         if context == "restart":
+            logging.info(f"Using latest model from epoch {checkpoint['epoch']}")
             model_state_dict = checkpoint["model_state_dict"]
-        elif context == "finetune" or context == "export":
+        elif context in {"finetune", "export"}:
+            logging.info(f"Using best model from epoch {checkpoint['best_epoch']}")
             model_state_dict = checkpoint["best_model_state_dict"]
-            if model_state_dict is None:
-                model_state_dict = checkpoint["model_state_dict"]
         else:
             raise ValueError("Unknown context tag for checkpoint loading!")
 
         # Create the model
+        model_data = checkpoint["model_data"]
         model = cls(
             hypers=model_data["model_hypers"],
             dataset_info=model_data["dataset_info"],
@@ -858,10 +858,11 @@ def _add_output(self, target_name: str, target: TargetInfo) -> None:
 
     @classmethod
     def upgrade_checkpoint(cls, checkpoint: Dict) -> Dict:
-        if checkpoint["model_ckpt_version"] == 1:
-            checkpoints.update_v1_v2(checkpoint["model_state_dict"])
-            checkpoints.update_v1_v2(checkpoint["best_model_state_dict"])
-            checkpoint["model_ckpt_version"] = 2
+        for v in range(1, cls.__checkpoint_version__):
+            if checkpoint["model_ckpt_version"] == v:
+                update = getattr(checkpoints, f"model_update_v{v}_v{v + 1}")
+                update(checkpoint)
+                checkpoint["model_ckpt_version"] = v + 1
 
         if checkpoint["model_ckpt_version"] != cls.__checkpoint_version__:
             raise RuntimeError(
@@ -880,8 +881,10 @@ def get_checkpoint(self) -> Dict:
                 "model_hypers": self.hypers,
                 "dataset_info": self.dataset_info,
             },
+            "epoch": None,
+            "best_epoch": None,
             "model_state_dict": self.state_dict(),
-            "best_model_state_dict": None,
+            "best_model_state_dict": self.state_dict(),
         }
         return checkpoint