Training

Basics

To train a model we set it up in training mode:

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

If we use training mode, the requested neural network is passed to the TrainingAdapter, which is a pl.LightningModule. The TrainingAdapterperforms the training. For hyperparameter search each of the implemented neural network potentials implements a _config_prior method that will return the hyperparameter search space. Calling model_trainer.get_ray_trainer() will return a configurated Ray Trainer that can be used to search for optimal hyperaprameters.