ksuit.callbacks.checkpoint_callbacks.ema
========================================

.. py:module:: ksuit.callbacks.checkpoint_callbacks.ema


Classes
-------

.. autoapisummary::

   ksuit.callbacks.checkpoint_callbacks.ema.EmaCallback


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

.. py:class:: EmaCallback(callback_config, **kwargs)

   Bases: :py:obj:`ksuit.callbacks.base.PeriodicCallback`


   Callback for exponential moving average (EMA) of model weights.

   Initializes the EmaCallback.

   :param callback_config: configuration of the `EmaCallback`. Implements the `CallBackBaseConfig` schema.
   :param \*\*kwargs: additional arguments passed to the parent class.


   .. py:attribute:: model_paths


   .. py:attribute:: target_factors


   .. py:attribute:: save_weights


   .. py:attribute:: save_last_weights


   .. py:attribute:: save_latest_weights


   .. py:attribute:: parameters
      :type:  dict[tuple[str | None, float], dict[str, torch.Tensor]]


   .. py:attribute:: buffers
      :type:  dict[str | None, dict[str, torch.Tensor]]


   .. py:method:: resume_from_checkpoint(resumption_paths, model)

      If a callback stores large files to disk and is stateful (e.g., an EMA of the model), it would be
      unecessarily wasteful to also store the state in the callbacks `state_dict`. Therefore,
      `resume_from_checkpoint` is called when resuming a run, which allows callbacks to load their state from any
      file that was stored on the disk.

      :param resumption_path: PathProvider instance to access paths from the checkpoint to resume from.
      :param model: model of the current training run.



   .. py:method:: apply_ema(cur_model, model_path, target_factor)

      fused in-place implementation