tff.learning.optimizers.schedule_learning_rate

Returns an optimizer with scheduled learning rate.

tff.learning.optimizers.schedule_learning_rate(
    optimizer: tff.learning.optimizers.Optimizer,
    schedule_fn: Callable[[optimizer_base.Int], optimizer_base.Float]
) -> tff.learning.optimizers.Optimizer

The returned optimizer will use a learning rate of schedule_fn(i) for the i-th invocation of its next method, indexing from 0.

optimizer A tff.learning.optimizers.Optimizer which uses a learning rate.
schedule_fn A callable mapping integer round number to a floating point learning rate. To be invoked in the context of a tff.tensorflow.computation, thus should support a tf.Tensor input.

A tff.learning.optimizers.Optimizer.

KeyError If the provided optimizer's state is not a dictionary with learning rate stored under the tff.learning.optimizers.LEARNING_RATE_KEY key.