helios.scheduler.schedulers
===========================

.. py:module:: helios.scheduler.schedulers


Classes
-------

.. autoapisummary::

   helios.scheduler.schedulers.CosineAnnealingRestartLR
   helios.scheduler.schedulers.LinearWarmupScheduler
   helios.scheduler.schedulers.MultiStepRestartLR


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

.. py:class:: CosineAnnealingRestartLR(optimizer: torch.optim.Optimizer, periods: list[int], restart_weights: list[float] | None = None, eta_min: float = 0, last_epoch: int = -1)

   Bases: :py:obj:`torch.optim.lr_scheduler.LRScheduler`


   A cosine annealing with restarts LR scheduler.

   .. rubric:: Example

   Given

   .. code-block:: text

       periods = [10, 10, 10, 10]
       restart_weights = [1, 0.5, 0.5, 0.5]
       eta_min = 1e-7

   Then the scheduler will have 4 cycles of 10 iterations each. At the 10th, 20th,
   and 30th, the scheduler will restart with the weights in ``restart_weights``.

   :param optimizer: the optimizer.
   :param periods: period for each cosine annealing cycle.
   :param restart_weights: (optional) restart weights at each restart iteration. Defaults
                           to ``[1]``.
   :param eta_min: (optional) the minimum lr.
   :param last_epoch: (optional) index of the last epoch seen by the scheduler.


   .. py:method:: get_lr() -> list[float | torch.Tensor]

      Return the current learning rate.



.. py:class:: LinearWarmupScheduler(optimizer: torch.optim.Optimizer, warmup_steps: int, scheduler: torch.optim.lr_scheduler.LRScheduler, warmup_start_factor: float = 0.0, last_epoch: int = -1)

   Bases: :py:obj:`torch.optim.lr_scheduler.LRScheduler`


   A linear warmup LR scheduler.

   The scheduler functions as follows:
       #. During the first ``warmup_steps``, increase the learning rate lineary from
       ``base_lr * warmup_start_factor`` to ``base_lr``.
       #. After ``warmup_steps``, switch to the provided ``scheduler``.

   :param optimizer: the optimizer.
   :param warmup_steps: number of warmup steps.
   :param scheduler: the post-warmup scheduler.
   :param warmup_start_factor: (optional) fraction of ``base_lr`` to start from.
   :param last_epoch: (optional) index of the last epoch seen by the scheduler.


   .. py:method:: get_lr() -> list[float | torch.Tensor]

      Return the current learning rate.



   .. py:method:: step(epoch: int | None = None) -> None

      Advance the scheduler by one step.



.. py:class:: MultiStepRestartLR(optimizer: torch.optim.Optimizer, milestones: list[int], gamma: float = 0.1, restarts: list[int] | None = None, restart_weights: list[float] | None = None, last_epoch: int = -1)

   Bases: :py:obj:`torch.optim.lr_scheduler.LRScheduler`


   Multi-step with restarts LR scheduler.

   :param optimizer: torch optimizer.
   :param milestones: iterations that will decrease learning rate.
   :param gamma: (optional) decrease ratio.
   :param restarts: (optional) restart iterations. Defaults to ``[0]``.
   :param restart_weights: (optional) restart weights at each restart iteration. Defaults
                           to ``[1]``.
   :param last_epoch: (optional) index of the last epoch seen by the scheduler.


   .. py:method:: get_lr() -> list[float | torch.Tensor]

      Return the current learning rate.