scaffold.data.artifact_manager.artifact_logger
==============================================

.. py:module:: scaffold.data.artifact_manager.artifact_logger


Attributes
----------

.. autoapisummary::

   scaffold.data.artifact_manager.artifact_logger.STATE_FILENAME
   scaffold.data.artifact_manager.artifact_logger.logger


Classes
-------

.. autoapisummary::

   scaffold.data.artifact_manager.artifact_logger.ModelLogger


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

.. py:class:: ModelLogger(artifact_manager: scaffold.data.artifact_manager.base.ArtifactManager)

   Class to be used for logging a model or state dicts as an artifact.
   Will always use a temporary dir locally to save artifacts before logging that dir.

   Initializes the ModelLogger.

   :param artifact_manager: ArtifactManager to use for logging artifacts.
   :type artifact_manager: ArtifactManager


   .. py:method:: log_state_to_artifact(afid: str, model: torch.nn.Module, artifact_description: str, optimizers: List[torch.optim.Optimizer] = None, collection: Optional[str] = None, **kwargs) -> str

      Logs the model and optimizer state dicts under the specified artifact id.

      :param afid: Artifact id to log under on the given artifact store or None to generate a new one.
      :type afid: str
      :param model: Model to save the state dict off.
      :type model: torch.nn.Module
      :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 optimizers: All optimizers to save the state dicts off.
      :type optimizers: List[torch.optim.Optimizer]
      :param collection: Collection to log the artifact to.
      :type collection: Optional[str]
      :param kwargs: Additional key value pairs to save to the state dict



   .. py:method:: retrieve_state_from_artifact(afid: str, collection: Optional[str] = None, version: int = None, device: Optional[str] = None) -> Dict

      Returns a dictionary containing the model and optimizer state from the artifact store at the artifact_path.

      :param afid: Artifact id of the model. Must have been logged via :py:meth:`log_state_to_artifact`
      :param device: Device to load the model to. "cpu", "cuda" or None, which is the default that
                     will use the gpu if available

      :returns: State dictionary containing the model and optimizer state. See :py:meth:`log_state_to_artifact` for the
                key-value pairs.



   .. py:method:: save_state(dict_dp: str, model: torch.nn.Module, optimizers: List[torch.optim.Optimizer] = None, schedulers: List[torch.optim.lr_scheduler._LRScheduler] = None, **kwargs) -> None
      :staticmethod:


      Saves the model and optimizer state dicts, and optionally some keyword arguments, into one dictionary
      that can later be loaded to retrieve the exact state of the model training.

      .. rubric:: Notes

      We intentionally do not save the model.state_dict() here since we need to support easy
      model loading from a single afid in the evaluator without initializing the model class first.

      :param dict_dp: Directory where the state dict pickle file should be saved.
      :type dict_dp: str
      :param model: Model to save the state dict off.
      :type model: torch.nn.Module
      :param optimizers: All optimizers to save the state dicts off.
      :type optimizers: List[torch.optim.Optimizer]
      :param kwargs: Additional key value pairs to save to the state dict



   .. py:attribute:: artifact_manager


.. py:data:: STATE_FILENAME
   :value: 'state.pt'


.. py:data:: logger