Wrapping architectures to ensure consistent behavior and reduce workload on devs and maintainers.

src/metatrain/cli/train.py

@@ -635,51 +635,47 @@ 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}'")
-            checkpoint = torch.load(
-                restart_from, weights_only=False, map_location="cpu"
-            )
-            try:
-                model = model_from_checkpoint(checkpoint, context="restart")
-            except Exception as e:
-                raise ValueError(
-                    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_hypers=new_model_hypers)
-            try:
-                trainer = trainer_from_checkpoint(
-                    checkpoint=checkpoint,
-                    hypers=hypers["training"],
-                    context=training_context,  # type: ignore
-                )
-            except Exception as e:
-                raise ValueError(
-                    f"The file {restart_from} does not contain a valid checkpoint for "
-                    f"the '{architecture_name}' trainer state"
-                ) from e
-        elif training_context == "finetune" and restart_from is not None:
-            logging.info(f"Starting finetuning from '{restart_from}'")
+        if restart_from is not None:
             checkpoint = torch.load(
                 restart_from, weights_only=False, map_location="cpu"
             )
+            new_model_hypers = input_options.get("architecture", {}).get("model", {})
+
+            # Initialize the trainer.
+            if training_context == "restart":
+                logging.info(f"Restarting training from '{restart_from}'")
+                try:
+                    trainer = trainer_from_checkpoint(
+                        checkpoint=checkpoint,
+                        hypers=hypers["training"],
+                        context=training_context,  # type: ignore
+                    )
+                except Exception as e:
+                    raise ValueError(
+                        f"The file {restart_from} does not contain a valid checkpoint for "
+                        f"the '{architecture_name}' trainer state"
+                    ) from e
+            else:
+                logging.info(f"Starting finetuning from '{restart_from}'")
+                trainer = Trainer(hypers["training"])
+
+            # Load the model from the checkpoint.
             try:
-                model = model_from_checkpoint(checkpoint, context="finetune")
+                model = model_from_checkpoint(checkpoint, context=training_context)
             except Exception as e:
                 raise ValueError(
                     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_hypers=new_model_hypers)
-            trainer = Trainer(hypers["training"])
+
+            # Make the trainer setup the model to continue training.
+            model = trainer.restart(model, dataset_info, model_hypers=new_model_hypers)
+
         else:
             logging.info("Starting training from scratch")
-            model = Model(hypers["model"], dataset_info)
             trainer = Trainer(hypers["training"])
+            model = trainer.setup(hypers["model"], dataset_info)
     except Exception as e:
         raise ArchitectureError(e) from e
 

src/metatrain/composition/model.py

@@ -594,8 +594,6 @@ def export(self, metadata: Optional[ModelMetadata] = None) -> AtomisticModel:
         :return: An instance of :py:class:`metatomic.torch.AtomisticModel`.
         """
         dtype = self.dummy_buffer.dtype
-        if dtype not in self.__supported_dtypes__:
-            raise ValueError(f"unsupported dtype {dtype} for composition model")
 
         self.to(dtype)
         self.weights_to(torch.device("cpu"), torch.float64)

src/metatrain/composition/trainer.py

@@ -29,10 +29,10 @@
 from metatrain.utils.transfer import batch_to
 
 from . import checkpoints
-from .documentation import TrainerHypers
+from .documentation import TrainerHypers, ModelHypers
 
 
-class Trainer(TrainerInterface[TrainerHypers]):
+class Trainer(TrainerInterface[TrainerHypers, ModelHypers]):
     __checkpoint_version__ = 2
 
     def __init__(self, hypers: TrainerHypers):

src/metatrain/pet/checkpoints.py

@@ -373,6 +373,25 @@ def model_update_v15_v16(checkpoint: dict) -> None:
                     updated[name] = value
             checkpoint[key] = updated
 
+def model_update_v16_v17(checkpoint: dict) -> None:
+    """
+    Update a v16 checkpoint to v17.
+
+    It removes the additive models and scaler from the model checkpoint,
+    as this is now handled by the MetatrainModel wrapper.
+
+    :param checkpoint: The checkpoint to update.
+    """
+    removed_prefixes = (
+        "additive_models.",
+        "scaler.",
+    )
+    for key in ["model_state_dict", "best_model_state_dict"]:
+        if (state_dict := checkpoint.get(key)) is not None:
+            for k in list(state_dict):
+                for prefix in removed_prefixes:
+                    if k.startswith(prefix):
+                        state_dict.pop(k)
 
 ###########################
 # TRAINER #################