deepmd.pd.utils.learning_rate

Classes

LearningRateExp

Exponential decay learning rate schedule with optional warmup.

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