kliff.trainer.utils.lightning_utils

class kliff.trainer.utils.lightning_utils.SaveModelCallback(*args: Any, **kwargs: Any)

Callback to save the model at the end of each epoch. The model is saved in the ckpt_dir with the name “last_model.pth”. The best model is saved with the name “best_model.pth”. The model is saved every ckpt_interval epochs with the name “epoch_{epoch}.pth”.

on_validation_epoch_end(trainer, pl_module)

class kliff.trainer.utils.lightning_utils.SavePerAtomPredictions(*args: Any, **kwargs: Any)

Callback to save the per atom predictions of the model during validation. The per atom predictions are saved in the supplied lmdb file. Usually it is named per_atom_pred_database.lmdb in the run dir

on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

Return type:

None