View source on GitHub
|
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.
Methods
initialize
@abc.abstractmethodinitialize( specs )
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
@abc.abstractmethodnext( state, weights, gradients )
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.
|