Training

Basics

We use Pytorch Lightning to train the implemented potentials. To obtain a model that can be trained with PyTorch Lightning, the keyword training has to be passed to the NeuralNetworkPotentialFactory.create_nnp.
This will tell the factory to initialize the requested model and pass it to the TrainginAdapter, which then uses Pytorch Lightning for its training.

from modelforge.potential import NeuralNetworkPotentialFactory
# Set up model
model_trainer = NeuralNetworkPotentialFactory.create_nnp("training", "ANI2x")

To train the model we then need to set up the trainer and fit the model on a training set.

# set up traininer
from lightning.pytorch.callbacks.early_stopping import EarlyStopping

trainer = Trainer(
    max_epochs=10_000,
    num_nodes=1,
    devices=1,
    accelerator="cpu",
    logger=logger,  # Add the logger here
    callbacks=[
        EarlyStopping(monitor="val_loss", min_delta=0.05, patience=20, verbose=True)
    ],
)


# Run training loop and validate
trainer.fit(model, dataset.train_dataloader(), dataset.val_dataloader())

A complete training example is provided in the training script.

Atomic self-energies

Energies obtained using quantum mechanics typically contain an atomic energy offset, often called atomic self-energies. These self-energies are the energies that remain if a system loses every interatomic interaction (e.g., by increasing the interatomic distances to infinity (or, more realistically, above the interaction cutoff)). It has been observed that removing this offset helps with learning rates. Removing the offset does not change the width of the dataset energy distribution, i.e. (E_min - E_max) == (E_min_offset - E_max_offset), but moves the distribution closer to zero.

Besides faster training convergence, this is also relevant for increasing numerical precision: typically, the training is performed with a single precision, and a large offset might limit the number of relevant bits. The offset is removed from the dataset --- and energies are provided in double precision. After removing the offset, the remaining energy is cast to a PyTorch tensor and single precision.

We suggest always training on energies that have the atomic self-energies removed. Modelforge provides methods to perform this during the setup of the dataset, and by default, self-energies are removed:

dataset.prepare_data(remove_self_energies=True, normalize=False)

The atomic self-energies depend on the QM method used to calculate the data points of the training set. Typically, these are provided for each element in each dataset. If this information is missing, modelforge will estimate the self-energies using linear regression.