scaffold.torch.lightning.callbacks
==================================

.. py:module:: scaffold.torch.lightning.callbacks


Attributes
----------

.. autoapisummary::

   scaffold.torch.lightning.callbacks.logger


Classes
-------

.. autoapisummary::

   scaffold.torch.lightning.callbacks.LightningCheckpointer


Module Contents
---------------

.. py:class:: LightningCheckpointer(artifact_manager: scaffold.data.artifact_manager.base.ArtifactManager, artifact_description: str, target_afid: Optional[str] = None, target_afid_best: Optional[str] = None, resume_checkpoint_afid: Optional[str] = None, resume_checkpoint_version: Optional[int] = None, only_log_current_best: bool = True)

   Bases: :py:obj:`pytorch_lightning.callbacks.Callback`


   
   Initializing the custom checkpointer.

   .. rubric:: Notes

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

   :param artifact_manager: ArtifactManager to use for logging.
   :type artifact_manager: ArtifactManager
   :param artifact_description: Description of the artifact.
                                Will be logged at <artifact_root>/ARTIFACT_META_DIR/ARTIFACT_DESCRIPTION_FILE
                                and serves to reduce undocumented artifact clutter.
   :type artifact_description: str
   :param target_afid: Afid to log the state of the model and optimizers at the end of every epoch.
   :type target_afid: str
   :param target_afid_best: If set, will save the best model with this afid name. If None, will create a new
                            random afid with the prefix 'best_model'
   :type target_afid_best: str
   :param resume_checkpoint_afid: Afid to load a model, optimizer and progress state from at the start of
                                  training.
   :type resume_checkpoint_afid: str
   :param resume_checkpoint_version: Artifact version to load from. If None, will load the latest version.
   :type resume_checkpoint_version: int
   :param only_log_current_best: 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.
   :type only_log_current_best: bool


   .. py:method:: load_best_state() -> dict

      Convenience function to load the current local best state.



   .. py:method:: on_train_end(trainer: pytorch_lightning.Trainer, pl_module: pytorch_lightning.LightningModule) -> None

      Log the best state dicts under the target afid



   .. py:method:: 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.



   .. py:method:: 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.



   .. py:attribute:: artifact_description


   .. py:attribute:: artifact_manager


   .. py:attribute:: best_state_dir


   .. py:attribute:: best_state_path


   .. py:attribute:: lowest_avg_val_loss
      :value: None



   .. py:attribute:: model_logger


   .. py:attribute:: only_log_current_best
      :value: True



   .. py:attribute:: resume_checkpoint_afid
      :value: None



   .. py:attribute:: resume_checkpoint_version
      :value: None



   .. py:attribute:: target_afid
      :value: None



   .. py:attribute:: target_afid_best
      :value: None



.. py:data:: logger