Check input model hypers when restarting trainings

src/metatrain/cli/train.py

@@ -298,6 +298,7 @@ def train_model(
     # MERGE OPTIONS ###########
     ###########################
 
+    input_options = copy.deepcopy(options)
     options = OmegaConf.merge(
         BASE_OPTIONS,
         {"architecture": get_default_hypers(architecture_name)},
@@ -631,6 +632,8 @@ def train_model(
     else:
         training_context = None
 
+    new_model_hypers = input_options.get("architecture", {}).get("model", {})
+
     try:
         if training_context == "restart" and restart_from is not None:
             logging.info(f"Restarting training from '{restart_from}'")
@@ -644,7 +647,7 @@ def train_model(
                     f"The file {restart_from} does not contain a valid checkpoint for "
                     f"the '{architecture_name}' architecture"
                 ) from e
-            model = model.restart(dataset_info)
+            model = model.restart(dataset_info, model_hypers=new_model_hypers)
             try:
                 trainer = trainer_from_checkpoint(
                     checkpoint=checkpoint,
@@ -668,7 +671,7 @@ def train_model(
                     f"The file {restart_from} does not contain a valid checkpoint for "
                     f"the '{architecture_name}' architecture"
                 ) from e
-            model = model.restart(dataset_info)
+            model = model.restart(dataset_info, model_hypers=new_model_hypers)
             trainer = Trainer(hypers["training"])
         else:
             logging.info("Starting training from scratch")

src/metatrain/composition/model.py

@@ -1,6 +1,6 @@
 import logging
 import warnings
-from typing import Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 import metatensor.torch as mts
 import torch
@@ -22,6 +22,7 @@
     sparsify_atomic_basis_target,
 )
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.metadata import merge_metadata
 
 from . import checkpoints
@@ -204,7 +205,9 @@ def train_model(
             checkpoint_dir="",
         )
 
-    def restart(self, dataset_info: DatasetInfo) -> "CompositionModel":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "CompositionModel":
         """
         Update the model to continue training, possibly with new targets.
 
@@ -214,8 +217,13 @@ def restart(self, dataset_info: DatasetInfo) -> "CompositionModel":
 
         :param dataset_info: Information about the new dataset, including the
             targets that will be used for training.
+        :param model_hypers: New hyperparameters for the model. They must match
+            the existing hyperparameters, otherwise an error is raised.
         :return: The updated model.
         """
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         raw_targets = {}
         for target_name in dataset_info.targets:
             target_info = dataset_info.targets[target_name]

src/metatrain/experimental/classifier/model.py

@@ -123,7 +123,9 @@ def build_mlp(self, feature_size: int, num_classes: int) -> None:
         # Final classification layer
         self.linear = torch.nn.Linear(current_size, num_classes, bias=False)
 
-    def restart(self, dataset_info: DatasetInfo) -> "Classifier":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "Classifier":
         raise ValueError("Restarting from a Classifier model is not supported.")
 
     def forward(

src/metatrain/experimental/dpa3/model.py

@@ -25,6 +25,7 @@
 from metatrain.utils.data.atom_pair_helpers import check_no_atom_pair_targets
 from metatrain.utils.data.dataset import DatasetInfo
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
 from metatrain.utils.sum_over_atoms import sum_over_atoms
@@ -391,7 +392,13 @@ def forward(
 
         return return_dict
 
-    def restart(self, dataset_info: DatasetInfo) -> "DPA3":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "DPA3":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/experimental/flashmd/model.py

@@ -28,6 +28,7 @@
 from metatrain.utils.data import DatasetInfo, TargetInfo
 from metatrain.utils.data.atom_pair_helpers import check_no_atom_pair_targets
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.long_range import DummyLongRangeFeaturizer, LongRangeFeaturizer
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
@@ -232,7 +233,13 @@ def set_timestep(self, timestep: float):
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "FlashMD":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "FlashMD":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/experimental/flashmd_symplectic/model.py

@@ -31,6 +31,7 @@
 from metatrain.utils.data.atom_pair_helpers import check_no_atom_pair_targets
 from metatrain.utils.data.target_info import get_energy_target_info
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.long_range import DummyLongRangeFeaturizer, LongRangeFeaturizer
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
@@ -222,7 +223,13 @@ def set_timestep(self, timestep: float):
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "FlashMDSymplectic":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "FlashMDSymplectic":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/experimental/mace/model.py

@@ -30,6 +30,7 @@
     sparsify_atomic_basis_target,
 )
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
 from metatrain.utils.sum_over_atoms import sum_over_atoms
@@ -306,7 +307,13 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
 
         self.finetune_config: Dict[str, Any] = {}
 
-    def restart(self, dataset_info: DatasetInfo) -> "MetaMACE":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "MetaMACE":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # Check that the new dataset info does not contain new atomic types
         if new_atomic_types := set(dataset_info.atomic_types) - set(
             self.dataset_info.atomic_types

src/metatrain/experimental/space/model.py

@@ -35,6 +35,7 @@
 )
 from metatrain.utils.data.dataset import DatasetInfo, TargetInfo
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
 
@@ -166,7 +167,13 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "SPACE":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "SPACE":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/experimental/space/tests/test_basic.py

@@ -11,6 +11,7 @@
 from metatrain.utils.testing import (
     ArchitectureTests,
     CheckpointTests,
+    InputTests,
     OutputTests,
     TorchscriptTests,
 )
@@ -31,6 +32,9 @@ def minimal_model_hypers(self):
         return hypers
 
 
+class TestInput(InputTests, SPACETests): ...
+
+
 class TestOutput(OutputTests, SPACETests):
     is_equivariant_reflections = False
     equivariance_error_tolerance = 1e-4  # due to many layers in the default hypers

src/metatrain/gap/model.py

@@ -169,7 +169,9 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "GAP":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "GAP":
         raise NotImplementedError("GAP does not allow restarting training")
 
     @classmethod

src/metatrain/llpr/model.py

@@ -26,6 +26,7 @@
 from metatrain.utils.data.target_info import (
     is_auxiliary_output,
 )
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.io import model_from_checkpoint
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.neighbor_lists import (
@@ -250,7 +251,13 @@ def set_wrapped_model(self, model: ModelInterface) -> None:
                 bias=False,
             )
 
-    def restart(self, dataset_info: DatasetInfo) -> "LLPRUncertaintyModel":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "LLPRUncertaintyModel":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/llpr/tests/test_basic.py

@@ -8,15 +8,24 @@
 from metatrain.pet import PET
 from metatrain.pet import Trainer as PETTrainer
 from metatrain.utils.architectures import get_default_hypers
+from metatrain.utils.data import DatasetInfo
 from metatrain.utils.hypers import init_with_defaults
 from metatrain.utils.loss import LossSpecification
 from metatrain.utils.testing import ArchitectureTests, CheckpointTests
+from metatrain.utils.testing.input import InputTests
 
 
 class LLPRTests(ArchitectureTests):
     architecture = "llpr"
 
 
+class TestInput(InputTests, LLPRTests):
+    def test_restart(
+        self, minimal_model_hypers: dict, dataset_info: DatasetInfo
+    ) -> None:
+        pytest.skip("LLPR's restart does not work without a model checkpoint.")
+
+
 class TestCheckpoints(CheckpointTests, LLPRTests):
     @pytest.fixture
     def model_trainer(

src/metatrain/pet/model.py

@@ -29,6 +29,7 @@
     sparsify_atomic_basis_target,
 )
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.long_range import DummyLongRangeFeaturizer, LongRangeFeaturizer
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
@@ -202,7 +203,13 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "PET":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "PET":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/soap_bpnn/model.py

@@ -27,6 +27,7 @@
 )
 from metatrain.utils.data.dataset import DatasetInfo
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.long_range import DummyLongRangeFeaturizer, LongRangeFeaturizer
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
@@ -423,7 +424,13 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "SoapBpnn":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "SoapBpnn":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/utils/abc.py

@@ -135,7 +135,9 @@ def supported_outputs(self) -> Dict[str, ModelOutput]:
         """
 
     @abstractmethod
-    def restart(self, dataset_info: DatasetInfo) -> "ModelInterface":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "ModelInterface":
         """
         Update a model to restart training, potentially with different dataset and/or
         targets.
@@ -147,6 +149,8 @@ def restart(self, dataset_info: DatasetInfo) -> "ModelInterface":
         :param dataset_info: Information about the new dataset, including the targets
             that will be used for training.
 
+        :param model_hypers: The new hyperparameters for the model.
+
         :return: The updated model, or a new instance of the model, that is able to
             handle the new dataset.
         """