Model I/O

mlip.models.model_io.save_model_to_zip(save_path: str | PathLike, model: ForceField) → None

Saves a force field model to a zip archive in a lightweight format to be easily loaded back for inference later.

Parameters:

  • save_path – The target path to the zip archive. Should have extension “.zip”.

  • model – The force field model to save. Must be passed as type ForceField.

mlip.models.model_io.load_model_from_zip(model_type: type, load_path: str | PathLike) → ForceField

Loads a model from a zip archive and returns it wrapped as a ForceField.

Parameters:

  • model_type – The model class that corresponds to the saved model.

  • load_path – The path to the zip archive to load.

Returns:

The loaded model wrapped as a ForceField object.

mlip.models.params_loading.load_parameters_from_checkpoint(local_checkpoint_dir: str | PathLike, initial_params: dict[str, dict[str, Array | dict]], epoch_to_load: int, load_ema_params: bool = False) → dict[str, dict[str, Array | dict]]

Loads model parameters from a checkpoint.

Parameters:

  • local_checkpoint_dir – The directory (must be local) where the checkpoints are stored. This directory should contain the subdirectories named after the epoch numbers of the checkpoints.

  • initial_params – The initial parameters of the model as a template for loading.

  • epoch_to_load – The epoch number to load.

  • load_ema_params – Whether to load the EMA parameters instead of the standard ones. By default, this is set to False.

Returns:

The loaded model parameters.