tff.learning.optimizers.Optimizer
Stay organized with collections
Save and categorize content based on your preferences.
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.TensorSpec s 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.
|
Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates.
Last updated 2024-09-20 UTC.
[[["Easy to understand","easyToUnderstand","thumb-up"],["Solved my problem","solvedMyProblem","thumb-up"],["Other","otherUp","thumb-up"]],[["Missing the information I need","missingTheInformationINeed","thumb-down"],["Too complicated / too many steps","tooComplicatedTooManySteps","thumb-down"],["Out of date","outOfDate","thumb-down"],["Samples / code issue","samplesCodeIssue","thumb-down"],["Other","otherDown","thumb-down"]],["Last updated 2024-09-20 UTC."],[],[],null,["# tff.learning.optimizers.Optimizer\n\n\u003cbr /\u003e\n\n|-------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| [View source on GitHub](https://github.com/tensorflow/federated/blob/v0.87.0 Version 2.0, January 2004 Licensed under the Apache License, Version 2.0 (the) |\n\nRepresents an optimizer for use in TensorFlow Federated.\n\nIts pair of `initialize` and `next` methods define the optimization\nalgorithm, with `next` corresponding to a step of the optimizer.\n\nThis class captures iterative optimization algorithms where the same operation\nis applied in every optimization step. The `next` method should be a\ncomputation that can be implemented as a [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function). Each method will\ngenerally only be traced once to create the corresponding TensorFlow graph\nfunctions. Thus, the methods should not create or use [`tf.Variable`](https://www.tensorflow.org/api_docs/python/tf/Variable) objects.\n\nInstead, any dependency between steps of the algorithm should be included\nas tensors in a state. For instance, a momentum term for momentum SGD is\ncreated in the `initialize` method as all-zeros tensor, which is then both\nan input and an output of the `next` method. These aspects should be accessed\nand changed via `get_hparams` and `set_hparams`, respectively.\n\nAs a best practice, any implementation using learning rate, should store it in\nits state under the key [`tff.learning.optimizers.LEARNING_RATE_KEY`](../../../tff/learning/optimizers#LEARNING_RATE_KEY).\n\nMethods\n-------\n\n### `get_hparams`\n\n[View source](https://github.com/tensorflow/federated/blob/v0.87.0\nVersion 2.0, January 2004\nLicensed under the Apache License, Version 2.0 (the) \n\n get_hparams(\n state: State\n ) -\u003e Hparams\n\nReturns a dictionary containing the optimizer state hyperparameters.\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ||\n|---------|-------------------------------------------------------------------------------------------|\n| `state` | The state of the optimizer. Must match the structure returned by the `initialize` method. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ||\n|---|---|\n| An ordered dictionary representing the hyperparameters in the given state. ||\n\n\u003cbr /\u003e\n\n### `initialize`\n\n[View source](https://github.com/tensorflow/federated/blob/v0.87.0\nVersion 2.0, January 2004\nLicensed under the Apache License, Version 2.0 (the) \n\n @abc.abstractmethod\n initialize(\n specs: Any\n ) -\u003e State\n\nReturns the initial state of the optimizer.\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ||\n|---------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| `specs` | A (possibly nested) structure of [`tf.TensorSpec`](https://www.tensorflow.org/api_docs/python/tf/TensorSpec)s describing the weights to be optimized. The `weights` and `grads` argument of `next` must match the structure and (shape, dtype) of `specs`. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ||\n|---|---|\n| Initial state of the optimizer. A (possibly nested) structure of tensors. ||\n\n\u003cbr /\u003e\n\n### `next`\n\n[View source](https://github.com/tensorflow/federated/blob/v0.87.0\nVersion 2.0, January 2004\nLicensed under the Apache License, Version 2.0 (the) \n\n @abc.abstractmethod\n next(\n state: State, weights: Weights, gradients: Any\n ) -\u003e tuple[State, Weights]\n\nTakes a step of the optimizer.\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ||\n|-------------|------------------------------------------------------------------------------------------------------------------------------------------------------|\n| `state` | State of the optimizer. A structure of tensors matching the structure returned by `initialize` method. |\n| `weights` | The weights to be updated by the optimizer. A collection of tensors matching the structure of `specs` provided in the `initialize` method. |\n| `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. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ||\n|---|---|\n| A (state, weights) tuple representing the updated `state` and `weights`. ||\n\n\u003cbr /\u003e\n\n### `set_hparams`\n\n[View source](https://github.com/tensorflow/federated/blob/v0.87.0\nVersion 2.0, January 2004\nLicensed under the Apache License, Version 2.0 (the) \n\n set_hparams(\n state: State, hparams: Hparams\n ) -\u003e State\n\nReturns an optimizer state with updated hyperparameters.\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ||\n|-----------|--------------------------------------------------------------------------------------------------|\n| `state` | The state of the optimizer. Must match the structure returned by the `initialize` method. |\n| `hparams` | A dictionary matching the output of `get_hparams` containing the updated hyperparameters to use. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ||\n|---|---|\n| An ordered dictionary representing the hyperparameters in the given state. ||\n\n\u003cbr /\u003e"]]