tff.learning.optimizers.Optimizer

Represents an optimizer for use in TensorFlow Federated.

Its pair of initialize and next methods define the optimization algorithm, with next corresponding to a step of the optimizer.

This class captures iterative optimization algorithms where the same operation is applied in every optimization step. The next method should be a computation that can be implemented as a tf.function. Each method will generally only be traced once to create the corresponding TensorFlow graph functions. Thus, the methods should not create or use tf.Variable objects.

Instead, any dependency between steps of the algorithm should be included as tensors in a state. For instance, a momentum term for momentum SGD is created in the initialize method as all-zeros tensor, which is then both an input and an output of the next method. These aspects should be accessed and changed via get_hparams and set_hparams, respectively.

As a best practice, any implementation using learning rate, should store it in its state under the key tff.learning.optimizers.LEARNING_RATE_KEY.

Methods

get_hparams

View source

get_hparams(
    state: State
) -> Hparams

Returns a dictionary containing the optimizer state hyperparameters.

Args
state The state of the optimizer. Must match the structure returned by the initialize method.

Returns
An ordered dictionary representing the hyperparameters in the given state.

initialize

View source

@abc.abstractmethod
initialize(
    specs: Any
) -> State

Returns the initial state of the optimizer.

Args
specs A (possibly nested) structure of tf.TensorSpecs describing the weights to be optimized. The weights and grads argument of next must match the structure and (shape, dtype) of specs.

Returns
Initial state of the optimizer. A (possibly nested) structure of tensors.

next

View source

@abc.abstractmethod
next(
    state: State, weights: Weights, gradients: Any
) -> tuple[State, Weights]

Takes a step of the optimizer.

Args
state State of the optimizer. A structure of tensors matching the structure returned by initialize method.
weights The weights to be updated by the optimizer. A collection of tensors matching the structure of specs provided in the initialize method.
gradients The gradients to use for the update by the optimizer. A collection of tensors matching the structure of specs provided in the initialize method.

Returns
A (state, weights) tuple representing the updated state and weights.

set_hparams

View source

set_hparams(
    state: State, hparams: Hparams
) -> State

Returns an optimizer state with updated hyperparameters.

Args
state The state of the optimizer. Must match the structure returned by the initialize method.
hparams A dictionary matching the output of get_hparams containing the updated hyperparameters to use.

Returns
An ordered dictionary representing the hyperparameters in the given state.