kliff.trainer.kim_trainer

class kliff.trainer.kim_trainer.KIMTrainer(training_manifest, model=None)

This class extends the base Trainer class for training OpenKIM physics based models. It will use the scipy optimizers. It will perform a check to exclude TorchML model driver based models, as they would be handled by Torch based trainers. It can read model tarballs as well as export the models as tarballs for ease of use. It will use the KIMModel class to load the model and set the parameters. It also provides explicit interface for parameters transformation.

Parameters:

training_manifest (dict) – The training_manifest dictionary.

setup_model()

Load either the installed KIM model, or install it from the source. If the model driver required is TorchML* family, then it will raise an error, as it should be handled by the DNNTrainer, or GNNLightningTrainer.

Path can be a folder containing the model, or a tar file. The model name is the KIM model name.

setup_optimizer()

Set up the optimizer based on the provided information. If the optimizer name is not provided, it will raise an error. It will use the ~:class:~scipy.optimize class for optimizers. It will raise an error if the optimizer is not supported.

loss(x)

Compute the loss function for the given parameters. It sets the KIM model parameters, compute the desired loss function doe all trainable properties and return the total loss after scaling losses with ~:class:~kliff.configuration.Weight.

Parameters:

x (np.ndarray) – The model parameters.

Returns:

The total loss.

Return type:

float

checkpoint(*args, **kwargs)

train_step(*args, **kwargs)

validation_step(*args, **kwargs)

get_optimizer(*args, **kwargs)

train()

Train the model using the provided optimizer. It will set the model parameters to the optimal values found by the optimizer. It will log the optimization status and the message. It will raise an error if the optimization fails.

save_kim_model()

Save the KIM model to the provided path. It will also generate a tarball if specified in the export manifest.