Model I/O

dipm.utils.model_io.save_model(save_path: str | PathLike, model: ForceFieldPredictor) → None

Saves a force field model to a safetensors file.

Parameters:

  • save_path – The target path to the model file. Should have extension “.safetensors”.

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

dipm.utils.model_io.load_model(load_path: str | PathLike, model_type: Type[ForceModel] | None = None, dtype: str | type[Any] | dtype | SupportsDType | Any | None = None, drop_force_head: bool = True, strict: bool = True, elements_to_drop: set[int] | None = None) → ForceFieldPredictor

Loads a model from a safetensors file and returns it wrapped as a ForceFieldPredictor.

Parameters:

  • load_path – The path to the safetensors file to load.

  • model_type (optional) – The model class that corresponds to the saved model. If you are using a known model listed in dipm.models.KNOWN_MODELS, this argument can be optional and inferred from the saved metadata. Otherwise, it must be provided.

  • dtype (optional) – The dtype in computations. If not provided, it will be the same as the the parameters dtype.

  • drop_force_head (optional) – If Ture, the head of the forces will be dropped if it exists in the saved model. Default is True because it is not recommended to use the head of the forces in inferece.

  • strict (optional) – If True, raises an error if the parameters name and shape are mismatched.

  • elements_to_drop (optional) – If provided, the elements in this set will be dropped from the loaded model.

Returns:

The loaded model wrapped as a ForceFieldPredictor object.