deepmd.pd.utils.learning_rate

Classes

LearningRateExp	Exponential decay learning rate schedule with optional warmup.
LearningRateWSD	Warmup-stable-decay learning rate schedule with configurable decay rules.

Module Contents

class deepmd.pd.utils.learning_rate.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: BaseLR

Exponential decay learning rate schedule with optional warmup.

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

Stepped mode (smooth=False, default):

\[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):

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

The learning rate decays continuously at every step.

where: - \(lr_0\) is start_lr (learning rate at the start of decay phase) - \(r\) is the decay rate decay_rate - \(t\) is the step index within the decay phase - \(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:

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

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

decay_steps = 5000

min_lr

smooth = False

_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:

stepint or Array

The step index relative to the end of warmup.

Returns:

Array

The decayed learning rate (absolute value).

class deepmd.pd.utils.learning_rate.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: BaseLR

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

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

Let \(\tau \in [0, 1]\) denote the normalized progress within the decay phase.

Inverse-linear mode (``decay_type=”inverse_linear”``):

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

Cosine mode (``decay_type=”cosine”``):

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

Linear mode (``decay_type=”linear”``):

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

decay_phase_ratio = 0.1

decay_type = 'inverse_linear'

decay_phase_steps

stable_steps

_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:

stepint or Array

The step index relative to the end of warmup.

Returns:

Array

The learning rate (absolute value).