Train

metatensor.models.cli.train.train_model(options: DictConfig | Dict, output: str = 'model.pt', checkpoint_dir: str | Path = '.', continue_from: str | None = None) → None[source]

Train an atomistic machine learning model using provided options.

This function sets up the dataset and model architecture, then runs the training process. The dataset is prepared by reading structural data and target values from specified paths. The model architecture is dynamically imported and instantiated based on the configuration. Training is executed with the specified hyperparameters, and the trained model is saved to a designated output path.

Parameters:

  • options (DictConfig | Dict) – DictConfig containing the training options

  • output (str) – Path to save the final model

  • checkpoint_dir (str | Path) – Path to save checkpoints and other intermediate output files like the fully expanded training options for a later restart.

  • continue_from (str | None) – File to continue training from.

Return type:

None