trainers.EnsembleOfModels

class EnsembleOfModels(trainers, ref_energy_fn_template=None)

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

Example

trainer_list = []
for i in range(4):
    trainer_list.append(trainers.ForceMatching(...))
trainer_ensemble = trainers.EnsembleOfModels(trainer_list)

trainer_ensemble.train(*args, **kwargs)
trained_params = trainer_ensemble.list_of_params

Methods

__init__(trainers[, ref_energy_fn_template])	A reference energy_fn_template can be provided, but is not mandatory due to the dependence of the template on the box via the displacement function.
checkpoint(name, object)	Marks attribute to be saved in a partial checkpoint.
load_energy_params(file_path)	Loads energy parameters.
move_to_device()	Move all attributes that are expected to be on device to device to avoid TracerExceptions after loading trainers from disk, i.e. loading numpy rather than device arrays.
restore(checkpoint)	Restores the trainer from a checkpoint.
save_energy_params(file_path[, save_format, ...])	Saves energy parameters.
save_trainer(save_path[, format])	Saves whole trainer, e.g. for production after training.
train(*args, **kwargs)	Training of any trainer should start by calling train.

Attributes

energy_fn

Returns the energy function for the current parameters.

list_of_params

params