deepmd.pd.utils.learning_rate
=============================

.. py:module:: deepmd.pd.utils.learning_rate


Classes
-------

.. autoapisummary::

   deepmd.pd.utils.learning_rate.LearningRateExp
   deepmd.pd.utils.learning_rate.LearningRateWSD


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

.. 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 !!


.. py:class:: LearningRateWSD(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, decay_phase_ratio: float = 0.1, decay_type: str = 'inverse_linear', **kwargs: Any)

   Bases: :py:obj:`BaseLR`


   
   Warmup-stable-decay learning rate schedule with configurable decay rules.

   The schedule uses the shared warmup implementation from :class:`BaseLR`,
   then keeps the learning rate at ``start_lr`` during the stable phase, and
   finally applies one of the supported decay rules.

   Let :math:`\tau \in [0, 1]` denote the normalized progress within the
   decay phase.

   **Inverse-linear mode (``decay_type="inverse_linear"``):**

   .. math::

       lr(t) = \frac{1}{
           \tau / lr_{\text{stop}} + (1 - \tau) / lr_0
       }

   **Cosine mode (``decay_type="cosine"``):**

   .. math::

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

   **Linear mode (``decay_type="linear"``):**

   .. math::

       lr(t) = lr_0 + \left(lr_{\text{stop}} - lr_0\right)\tau















   ..
       !! processed by numpydoc !!

   .. py:attribute:: decay_phase_ratio
      :value: 0.1



   .. py:attribute:: decay_type
      :value: 'inverse_linear'



   .. py:attribute:: decay_phase_steps


   .. py:attribute:: stable_steps


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

      
      Get the warmup-stable-decay 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 learning rate (absolute value).











      ..
          !! processed by numpydoc !!