deepmd.pt.utils.learning_rate
=============================

.. py:module:: deepmd.pt.utils.learning_rate


Classes
-------

.. autoapisummary::

   deepmd.pt.utils.learning_rate.BaseLR
   deepmd.pt.utils.learning_rate.LearningRateCosine
   deepmd.pt.utils.learning_rate.LearningRateExp


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

.. py:class:: BaseLR(start_lr: float, num_steps: int, stop_lr: float | None = None, stop_lr_ratio: float | None = None, warmup_steps: int = 0, warmup_ratio: float | None = None, warmup_start_factor: float = 0.0, **kwargs: Any)

   Bases: :py:obj:`abc.ABC`, :py:obj:`deepmd.utils.plugin.PluginVariant`, :py:obj:`make_plugin_registry`\ (\ :py:obj:`'lr'`\ )


   
   Helper class that provides a standard way to create an ABC using
   inheritance.
















   ..
       !! processed by numpydoc !!

   .. py:attribute:: warmup_start_lr


   .. py:attribute:: _start_lr


   .. py:attribute:: num_steps


   .. py:attribute:: decay_num_steps


   .. py:property:: start_lr
      :type: float


      
      Get the starting learning rate.





      :Returns:

          :class:`python:float`
              The starting learning rate.











      ..
          !! processed by numpydoc !!


   .. py:method:: _decay_value(step: int | deepmd.dpmodel.array_api.Array) -> deepmd.dpmodel.array_api.Array
      :abstractmethod:


      
      Get the decayed learning rate at the given step (after warmup).

      This method should implement the actual decay logic (exp, cosine, etc.)
      without considering warmup.

      :Parameters:

          **step** : :class:`python:int` or :obj:`Array`
              The step index relative to the end of warmup.
              For example, if warmup_steps=100 and total_step=150, this method
              will be called with step=50.



      :Returns:

          :obj:`Array`
              The decayed learning rate (absolute value, not factor).











      ..
          !! processed by numpydoc !!


   .. py:method:: value(step: int | deepmd.dpmodel.array_api.Array) -> deepmd.dpmodel.array_api.Array | float

      
      Get the learning rate at the given step, including warmup.


      :Parameters:

          **step** : :class:`python:int` or :obj:`Array`
              The absolute step index from the start of training.



      :Returns:

          :obj:`Array`
              The learning rate at the given step.











      ..
          !! processed by numpydoc !!


.. py:class:: LearningRateCosine(start_lr: float, num_steps: int, stop_lr: float | None = None, stop_lr_ratio: float | None = None, warmup_steps: int = 0, warmup_ratio: float | None = None, warmup_start_factor: float = 0.0, **kwargs: Any)

   Bases: :py:obj:`BaseLR`


   
   Cosine annealing learning rate schedule with optional warmup.

   The decay phase (after warmup) follows the cosine annealing formula:

   .. math::

       lr(t) = lr_{\text{stop}} + \frac{lr_0 - lr_{\text{stop}}}{2} \left(1 + \cos\left(\pi \frac{t}{T}\right)\right)

   where:
   - :math:`lr_0` is ``start_lr`` (learning rate at the start of decay phase)
   - :math:`lr_{\text{stop}}` is ``stop_lr`` (minimum learning rate)
   - :math:`t` is the step index within the decay phase
   - :math:`T = \text{num\_steps} - \text{warmup\_steps}` is the total
     number of decay steps

   Equivalently, using :math:`\alpha = lr_{\text{stop}} / lr_0`:

   .. math::

       lr(t) = lr_0 \cdot \left[\alpha + \frac{1}{2}(1 - \alpha) \left(1 + \cos\left(\pi \frac{t}{T}\right)\right)\right]















   ..
       !! processed by numpydoc !!

   .. py:attribute:: lr_min_factor


   .. py:method:: _decay_value(step: int | deepmd.dpmodel.array_api.Array) -> deepmd.dpmodel.array_api.Array

      
      Get the cosine-annealed learning rate at the given step.


      :Parameters:

          **step** : :class:`python:int` or :obj:`Array`
              The step index relative to the end of warmup.



      :Returns:

          :obj:`Array`
              The annealed learning rate (absolute value).











      ..
          !! processed by numpydoc !!


.. py:class:: LearningRateExp(start_lr: float, num_steps: int, stop_lr: float | None = None, stop_lr_ratio: float | None = None, decay_steps: int = 5000, decay_rate: float | None = None, warmup_steps: int = 0, warmup_ratio: float | None = None, warmup_start_factor: float = 0.0, smooth: bool = False, **kwargs: Any)

   Bases: :py:obj:`BaseLR`


   
   Exponential decay learning rate schedule with optional warmup.

   The decay phase (after warmup) follows the exponential decay formula.

   **Stepped mode (smooth=False, default):**

   .. math::

       lr(t) = lr_0 \cdot r^{\lfloor t / s \rfloor}

   The learning rate decays every ``decay_steps`` steps, creating a staircase
   pattern.

   **Smooth mode (smooth=True):**

   .. math::

       lr(t) = lr_0 \cdot r^{t / s}

   The learning rate decays continuously at every step.

   where:
   - :math:`lr_0` is ``start_lr`` (learning rate at the start of decay phase)
   - :math:`r` is the decay rate ``decay_rate``
   - :math:`t` is the step index within the decay phase
   - :math:`s` is ``decay_steps`` (the decay period)

   The decay rate is automatically computed from ``start_lr`` and ``stop_lr``
   over the total decay steps unless explicitly provided:

   .. math::

       r = \left(\frac{lr_{\text{stop}}}{lr_0}\right)^{\frac{s}{T}}

   where :math:`T = \text{num\_steps} - \text{warmup\_steps}` is the total
   number of decay steps, and :math:`lr_{\text{stop}}` is ``stop_lr``.















   ..
       !! processed by numpydoc !!

   .. py:attribute:: decay_steps
      :value: 5000



   .. py:attribute:: min_lr


   .. py:attribute:: smooth
      :value: False



   .. py:method:: _decay_value(step: int | deepmd.dpmodel.array_api.Array) -> deepmd.dpmodel.array_api.Array

      
      Get the exponential-decayed learning rate factor at the given step.


      :Parameters:

          **step** : :class:`python:int` or :obj:`Array`
              The step index relative to the end of warmup.



      :Returns:

          :obj:`Array`
              The decayed learning rate (absolute value).











      ..
          !! processed by numpydoc !!