trainers

The trainers module provides a high-level interface to train models via chemtrain’s algorithms. The Base Trainers provide the core algorithms, while Combining Trainers, allows the construction of more advanced training schemes from multiple algorithms.

Base Trainers

Difftre(init_params, optimizer[, ...])	Trainer class for parametrizing potentials via the DiffTRe method.
DifftreParallel(key, init_params, optimizer, ...)	Trainer class for parametrizing potentials via the DiffTRe method.
PropertyPrediction(error_fn, ...[, ...])	Trainer for direct prediction of molecular properties.
ForceMatching(init_params, optimizer, ...[, ...])	Parametrizes potential models via the Force Matching method.
RelativeEntropy(init_params, optimizer[, ...])	Trainer for relative entropy minimization.
SGMCForceMatching(sgmc_solver, init_samples)	Trainer for stochastic gradient Markov-chain Monte Carlo training based on force-matching.

Combining Trainers

InterleaveTrainers([sequential, ...])

Interleaves updates to train models using multiple algorithms.

Train Multiple Models

EnsembleOfModels(trainers[, ...])

Train an ensemble of models by starting optimization from different initial parameter sets, for use in uncertainty quantification applications.

Trainer Templates

MLETrainerTemplate(optimizer, init_state, ...)	Abstract class implementing common properties and methods of single point estimate Trainers using optax optimizers.
DataParallelTrainer(loss_fn, model, ...[, ...])	Trainer for parallelized MLE training based on a dataset.
EarlyStopping(params, criterion[, ...])	A class that saves the best parameter obtained so far based on the validation loss and determines whether the optimization can be stopped based on some stopping criterion.

Extensions

chemtrain provides some extensions to the Base Trainers, e.g., to log data to services allowing live tracking of the training process.

Logging to Weights and Biases

wandb_log_difftre(run, trainer[, plot_fns])

Logs DiffTRe training statistics to Weights & Biases.