deepmd.tf.utils.learning_rate

Module Contents

Classes

LearningRateExp

The exponentially decaying learning rate.

class deepmd.tf.utils.learning_rate.LearningRateExp(start_lr: float, stop_lr: float = 5e-08, decay_steps: int = 5000, decay_rate: float = 0.95)

The exponentially decaying learning rate.

The learning rate at step \(t\) is given by

\[\alpha(t) = \alpha_0 \lambda ^ { t / \tau }\]

where \(\alpha\) is the learning rate, \(\alpha_0\) is the starting learning rate, \(\lambda\) is the decay rate, and \(\tau\) is the decay steps.

Parameters:

start_lr

Starting learning rate \(\alpha_0\)

stop_lr

Stop learning rate \(\alpha_1\)

decay_steps

Learning rate decay every this number of steps \(\tau\)

decay_rate

The decay rate \(\lambda\). If stop_step is provided in build, then it will be determined automatically and overwritten.

build(global_step: deepmd.tf.env.tf.Tensor, stop_step: int | None = None) → deepmd.tf.env.tf.Tensor

Build the learning rate.

Parameters:

global_step

The tf Tensor prividing the global training step

stop_step

The stop step. If provided, the decay_rate will be determined automatically and overwritten.

Returns:

learning_rate

The learning rate

start_lr() → float

Get the start lr.

value(step: int) → float

Get the lr at a certain step.