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.

Additionally, this function saves the current library version into the zip.

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, required_properties: Properties | None = None, inference_context: InferenceContext | None = None) → 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.

  • required_properties – The properties required from the loaded model. Default is None which means that a default Properties object will be used (i.e. energy and forces). Set explicitly if you require stress, hessians, or other outputs.

  • inference_context – Optional context to apply to the loaded force field for inference. For MoE models, experts are contracted for this context before the force field is returned.

Returns:

The loaded model wrapped as a ForceField object.

mlip.models.params_loading.load_parameters_from_checkpoint(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:

  • checkpoint_dir – The directory (Orbax-compatible) where the model 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.