View source on GitHub
|
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.
Args | |
|---|---|
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.
|
Returns | |
|---|---|
A tff.learning.optimizers.Optimizer.
|
Raises | |
|---|---|
KeyError
|
If the provided optimizer's state is not a dictionary with
learning rate stored under the tff.learning.optimizers.LEARNING_RATE_KEY
key.
|