Abstract base classes

These are classes that define the interfaces that should be followed by metatrain components (models, trainers, etc.).

class metatrain.utils.abc.ModelInterface(hypers: ModelHypersType, dataset_info: DatasetInfo, metadata: ModelMetadata)

Bases: Module, Generic[ModelHypersType]

Abstract base class for a machine learning model in metatrain.

All architectures in metatrain must be implemented as sub-class of this class, and implement the corresponding methods.

Parameters:

• hypers (ModelHypersType) – A dictionary with the model’s hyper-parameters.

• dataset_info (DatasetInfo) – Information containing details about the dataset, such as target quantities and atomic types.

• metadata (ModelMetadata) – Metadata about the model, e.g. author, description, and references.

__checkpoint_version__: int

The current version of the model’s checkpoint.

This is used to upgrade checkpoints produced with earlier versions of the code. See Checkpoint versioning for more information.

__supported_devices__: List[device]

List of torch devices supported by this model architecture.

They should be sorted in order of preference since metatrain will use this and __supported_dtypes__ to determine, based on the user request and machines’ availability, the optimal dtype and device for training.

__supported_dtypes__: List[dtype]

List of torch dtypes supported by this model architecture.

They should be sorted in order of preference since metatrain will use this and __supported_devices__ to determine, based on the user request and machines’ availability, the optimal dtype and device for training.

__default_metadata__: ModelMetadata

Default metadata for this model architecture.

Can be used to provide references that will be stored in the exported model. The references are stored in a dictionary with keys implementation and architecture. The implementation key should contain references to the software used in the implementation of the architecture, while the architecture key should contain references about the general architecture.

hypers

The model hyper passed at initialization

dataset_info

The dataset info passed at initialization

metadata

The metadata passed at initialization

abstractmethod forward(systems: List[System], outputs: Dict[str, ModelOutput], selected_atoms: Labels | None = None) → Dict[str, TensorMap]

Execute the model for the given systems, computing the requested outputs.

Parameters:

• systems (List[System]) – List of systems to evaluate the model on.

• outputs (Dict[str, ModelOutput]) – Dictionary of outputs that the model should compute.

• selected_atoms (Labels | None) – Optional Labels specifying a subset of atoms to compute the outputs for. If None, the outputs are computed for all atoms in each system.

Returns:

A dictionary mapping each requested output name to the corresponding TensorMap containing the computed values.

Return type:

Dict[str, TensorMap]

See also

metatomic.torch.ModelInterface for more explanation about the different arguments.

abstractmethod supported_outputs() → Dict[str, ModelOutput]

Get the outputs currently supported by this model.

This will likely be the same outputs that are set as this model capabilities in ModelInterface.export().

Returns:

A dictionary of the supported outputs by this model.

Return type:

Dict[str, ModelOutput]

abstractmethod restart(dataset_info: DatasetInfo, model_hypers: Mapping[str, Any] | None = None) → ModelInterface

Update a model to restart training, potentially with different dataset and/or targets.

This function is called whenever training restarts, with the same or a different dataset. It enables transfer learning (changing the targets), and fine-tuning (same targets, different datasets)

Parameters:

• dataset_info (DatasetInfo) – Information about the new dataset, including the targets that will be used for training.

• model_hypers (Mapping[str, Any] | None) – The new hyperparameters for the model.

Returns:

The updated model, or a new instance of the model, that is able to handle the new dataset.

Return type:

ModelInterface

remove_output(target_name: str) → None

Remove a previously registered output target from the model.

Used to drop targets whose heads are no longer meaningful, e.g. after a backbone-altering fine-tuning run (full/lora) makes them stale.

Not every architecture supports this; the default implementation raises NotImplementedError.

Parameters:

target_name (str) – Name of the target to remove.

Return type:

None

abstractmethod classmethod load_checkpoint(checkpoint: Dict[str, Any], context: Literal['restart', 'finetune', 'export']) → ModelInterface

Create a model from a checkpoint (i.e. state dictionary).

Parameters:

• checkpoint (Dict[str, Any]) – Checkpoint’s state dictionary.

• context (Literal['restart', 'finetune', 'export']) – Context in which to load the model. Possible values are "restart" when restarting a stopped traininf run, "finetune" when loading a model for further fine-tuning or transfer learning, and "export" when loading a model for final export. When multiple checkpoints are stored together, this can be used to pick one of them depending on the context.

Returns:

An instance of the model.

Return type:

ModelInterface

abstractmethod export(metadata: ModelMetadata | None = None) → AtomisticModel

Turn this model into an instance of metatomic.torch.MetatensorAtomisticModel, containing the model itself, a definition of the model capabilities and some metadata about the model.

Parameters:

metadata (ModelMetadata | None) – additional metadata to add in the model as specified by the user.

Returns:

An instance of metatomic.torch.MetatensorAtomisticModel

Return type:

AtomisticModel

abstractmethod classmethod upgrade_checkpoint(checkpoint: Dict[str, Any], version: int | None = None) → Dict[str, Any]

Upgrade the checkpoint to the current version of the model.

Parameters:

• checkpoint (Dict[str, Any]) – Checkpoint’s state dictionary.

• version (int | None) – Version to which the checkpoint should be upgraded. If None, the checkpoint will be upgraded to the current version of the model.

Raises:

RuntimeError – if the checkpoint cannot be upgraded to the desired version.

Returns:

The upgraded checkpoint.

Return type:

Dict[str, Any]

abstractmethod get_checkpoint(best_model_state_dict: dict[str, Any] | None = None) → Dict[str, Any]

Get the checkpoint of the model. This should contain all the information needed by load_checkpoint to recreate the same model instance.

Parameters:

best_model_state_dict (dict[str, Any] | None) – Optional state dictionary of the best model (if available). If provided, it should likely be included in the checkpoint to allow for resuming training or evaluation from the best model state.

Returns:

The model’s checkpoint.

Return type:

Dict[str, Any]

class metatrain.utils.abc.TrainerInterface(hypers: TrainerHypersType)

Bases: Generic[TrainerHypersType, ModelHypersType]

Abstract base class for a model trainer in metatrain.

All architectures in metatrain must implement such a trainer, which is responsible for training the model. The trainer must be a be sub-class of this class, and implement the corresponding methods.

Parameters:

hypers (TrainerHypersType) – A dictionary with the trainer’s hyper-parameters.

__checkpoint_version__: int

The current version of the trainer’s checkpoint.

This is used to upgrade checkpoints produced with earlier versions of the code. See Checkpoint versioning for more information.

hypers

The trainer hypers passed at intialization

abstractmethod setup(model_hypers: ModelHypersType, dataset_info: DatasetInfo) → Any

Setup the trainer and return an initialized model ready for training.

Parameters:

• model_hypers (ModelHypersType) – The hyper-parameters of the model to be trained.

• dataset_info (DatasetInfo) – Information about the dataset to be used for training.

Returns:

An initialized MetatrainModel ready for training.

Return type:

Any

abstractmethod restart(model: Any, dataset_info: DatasetInfo, model_hypers: ModelHypersType) → Any

Restart a model for training with a new dataset info.

Parameters:

• model (Any) – The model to be restarted.

• dataset_info (DatasetInfo) – Information about the new dataset to be used for training.

• model_hypers (ModelHypersType) – The new hyperparameters to set for the model.

Returns:

The restarted model, ready for training with the new dataset info.

Return type:

Any

abstractmethod train(model: Any, dtype: dtype, devices: List[device], train_datasets: List[Dataset | Subset], val_datasets: List[Dataset | Subset], checkpoint_dir: str) → None

Train the model using the train_datasets. How to train the model is left to this class, using the hyper-parameter given in __init__.

Parameters:

• model (Any) – the model to train

• dtype (dtype) – torch.dtype used by the data in the datasets

• devices (List[device]) – torch.device to use for training the model. When training with more than one device (e.g. multi-GPU training), this can contains multiple devices.

• train_datasets (List[Dataset | Subset]) – datasets to use to train the model

• val_datasets (List[Dataset | Subset]) – datasets to use for model validation

• checkpoint_dir (str) – directory where checkpoints shoudl be saved

Return type:

None

abstractmethod save_checkpoint(model: ModelInterface, path: str | Path) → None

Save a checkoint of both the model and trainer state to the given path

Parameters:

• model (ModelInterface) – The model to save in the checkpoint.

• path (str | Path) – The path where to save the checkpoint.

Return type:

None

abstractmethod classmethod upgrade_checkpoint(checkpoint: dict[str, Any], version: int | None = None) → dict[str, Any]

Upgrade the checkpoint to the current version of the trainer.

Parameters:

• checkpoint (dict[str, Any]) – Checkpoint’s state dictionary.

• version (int | None) – Version to which the checkpoint should be upgraded. If None, the checkpoint should be upgraded to the current version of the trainer.

Raises:

RuntimeError – if the checkpoint cannot be upgraded to the current version of the trainer.

Returns:

The upgraded checkpoint.

Return type:

dict[str, Any]

abstractmethod classmethod load_checkpoint(checkpoint: Dict[str, Any], hypers: TrainerHypersType, context: Literal['restart', 'finetune']) → TrainerInterface

Create a trainer instance from data stored in the checkpoint.

Parameters:

• checkpoint (Dict[str, Any]) – Checkpoint’s state dictionary.

• hypers (TrainerHypersType) – Hyper-parameters for the trainer, as specified by the user.

• context (Literal['restart', 'finetune']) – Context in which to load the model. Possible values are "restart" when restarting a stopped traininf run, and "finetune" when loading a model for further fine-tuning or transfer learning. When multiple checkpoints are stored together, this can be used to pick one of them depending on the context.

Returns:

The loaded trainer instance.

Return type:

TrainerInterface

metatrain.utils.abc.common_upgrade_checkpoint(checkpoint: dict[str, Any], version: int | None, current_version: int, checkpoint_type: Literal['model', 'trainer'], checkpoints: ModuleType) → dict[str, Any]

Generic function to apply checkpoint upgrades for both models and trainers.

The models and trainers can use it if their checkpoint upgrade logic follows the same pattern of having a module with upgrade functions named as model_update_v{v}_v{v + 1} and trainer_update_v{v}_v{v + 1}, where v is the version number.

Otherwise, they can implement their own upgrade logic in their upgrade_checkpoint method, they are not forced to use this function.

Parameters:

• checkpoint (dict[str, Any]) – Checkpoint’s dictionary to upgrade.

• version (int | None) – Version to which the checkpoint should be upgraded. If None, the checkpoint will be upgraded to the current version of the model/trainer.

• current_version (int) – The current version of the model/trainer.

• checkpoint_type (Literal['model', 'trainer']) – Whether the checkpoint upgrade corresponds to a model or a trainer.

• checkpoints (ModuleType) – The module containing the upgrade functions for the model/trainer.

Raises:

RuntimeError – if the checkpoint cannot be upgraded to the current version of the model/trainer.

Returns:

The upgraded checkpoint.

Return type:

dict[str, Any]