scaffold.torch.lightning.callbacks

Attributes

logger

Classes

LightningCheckpointer

Initializing the custom checkpointer.

Module Contents

class scaffold.torch.lightning.callbacks.LightningCheckpointer(artifact_manager: scaffold.data.artifact_manager.base.ArtifactManager, artifact_description: str, target_afid: str | None = None, target_afid_best: str | None = None, resume_checkpoint_afid: str | None = None, resume_checkpoint_version: int | None = None, only_log_current_best: bool = True)

Bases: pytorch_lightning.callbacks.Callback

Initializing the custom checkpointer.

Notes

This method will create a local directory under /tmp (or equivalent depending on the OS) to save the best model checkpoint.

Parameters:

  • artifact_manager (ArtifactManager) – ArtifactManager to use for logging.

  • artifact_description (str) – Description of the artifact. Will be logged at <artifact_root>/ARTIFACT_META_DIR/ARTIFACT_DESCRIPTION_FILE and serves to reduce undocumented artifact clutter.

  • target_afid (str) – Afid to log the state of the model and optimizers at the end of every epoch.

  • target_afid_best (str) – If set, will save the best model with this afid name. If None, will create a new random afid with the prefix ‘best_model’

  • resume_checkpoint_afid (str) – Afid to load a model, optimizer and progress state from at the start of training.

  • resume_checkpoint_version (int) – Artifact version to load from. If None, will load the latest version.

  • only_log_current_best (bool) – If True, will only log a state, if the validation loss of an epoch is the lowest recorded one until then. If False, a state will be logged every epoch.

load_best_state() → dict

Convenience function to load the current local best state.

on_train_end(trainer: pytorch_lightning.Trainer, pl_module: pytorch_lightning.LightningModule) → None

Log the best state dicts under the target afid

on_train_start(trainer: pytorch_lightning.Trainer, pl_module: pytorch_lightning.LightningModule) → None

If an afid for a checkpoint was given, load the full state into the trainer and lightning module.

on_validation_epoch_end(trainer: pytorch_lightning.Trainer, pl_module: pytorch_lightning.LightningModule) → None

Saving state dicts of models and optimizers after every validation run.

artifact_description
artifact_manager
best_state_dir
best_state_path
lowest_avg_val_loss = None
model_logger
only_log_current_best = True
resume_checkpoint_afid = None
resume_checkpoint_version = None
target_afid = None
target_afid_best = None
scaffold.torch.lightning.callbacks.logger