diffusion_models.utils.schemas
==============================

.. py:module:: diffusion_models.utils.schemas




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

.. py:class:: TrainingConfiguration

   A training configuration for simple experiment management.


   .. py:attribute:: training_name
      :type:  str

      The name of the training.


   .. py:attribute:: batch_size
      :type:  int

      The batch size used for training.


   .. py:attribute:: learning_rate
      :type:  float

      The learning rate of the training.


   .. py:attribute:: number_of_epochs
      :type:  int

      The number of epoch used for training.


   .. py:attribute:: checkpoint_rate
      :type:  int
      :value: 100


      The rate at which checkpoints are saved..


   .. py:attribute:: mixed_precision_training
      :type:  bool
      :value: False


      Whether or not to use automatic mixed precision training.


   .. py:attribute:: gradient_clip
      :type:  Optional[float]
      :value: None


      Whether or not to clip gradients.


.. py:class:: LogConfiguration

   An object to manage logging configuration.


   .. py:attribute:: log_rate
      :type:  int
      :value: 10


      The rate at which training metrics are logged.


   .. py:attribute:: image_rate
      :type:  int
      :value: 50


      The rate at which images are generated for visualization. This can be used to validate model performance.


   .. py:attribute:: number_of_images
      :type:  int
      :value: 5


      The number of images to generate.


.. py:class:: BetaSchedulerConfiguration

   A simplified beta scheduler configuration.


   .. py:attribute:: steps
      :type:  int

      The number of steps in the beta scheduler.


   .. py:attribute:: betas
      :type:  torch.Tensor

      The beta values.


   .. py:attribute:: alpha_bars
      :type:  torch.Tensor

      The alpha bar values.


.. py:class:: Checkpoint

   A simplified checkpoint framework for easy saving and loading.


   .. py:attribute:: epoch
      :type:  int

      The current epoch.


   .. py:attribute:: model_state_dict
      :type:  Dict[str, Any]

      The model state dict.


   .. py:attribute:: optimizer_state_dict
      :type:  Dict[str, Any]

      The optimizer state dict.


   .. py:attribute:: scaler
      :type:  Optional[torch.cuda.amp.GradScaler]

      The GradScaler instance.


   .. py:attribute:: beta_scheduler_config
      :type:  BetaSchedulerConfiguration

      The beta scheduler configuration.


   .. py:attribute:: tensorboard_run_name
      :type:  Optional[str]
      :value: None


      The name of the tensorboard run.


   .. py:attribute:: image_channels
      :type:  int
      :value: 3


      The number of image channels used in the training.


   .. py:attribute:: loss
      :type:  Optional[float]
      :value: None


      The final loss value recorded.

      .. note:: This is a legacy parameter and will be removed in a future release.


   .. py:method:: from_file(file_path)
      :classmethod:


      Load and instantiate a checkpoint from a file.

      :param file_path: The path to the checkpoint file.

      :returns: A checkpoint instance.



   .. py:method:: to_file(file_path)

      Saves a checkpoint to a file.



.. py:class:: OldCheckpoint

   .. py:attribute:: epoch
      :type:  int


   .. py:attribute:: model_state_dict
      :type:  Dict[str, Any]


   .. py:attribute:: optimizer_state_dict
      :type:  Dict[str, Any]


   .. py:attribute:: scaler
      :type:  Optional[torch.cuda.amp.GradScaler]


   .. py:attribute:: tensorboard_run_name
      :type:  Optional[str]
      :value: None



   .. py:attribute:: loss
      :type:  Optional[float]
      :value: None



   .. py:method:: from_file(file_path)
      :classmethod:



   .. py:method:: to_file(file_path)


   .. py:method:: to_new_checkpoint(beta_scheduler)