tff.learning.templates.build_apply_optimizer_finalizer

Builds finalizer that applies a step of an optimizer.

tff.learning.templates.build_apply_optimizer_finalizer(
    optimizer_fn: Union[tff.learning.optimizers.Optimizer, Callable[[], tf.keras.optimizers.Optimizer]],
    model_weights_type: tff.types.StructType
)

Used in the notebooks

Used in the tutorials

The provided model_weights_type must be a non-federated tff.Type with the tff.learning.models.ModelWeights container.

The 2nd input argument of the created FinalizerProcess.next expects a value matching model_weights_type and its 3rd argument expects value matching model_weights_type.trainable. The optimizer will be applied to the trainable model weights only, leaving non_trainable weights unmodified.

The state of the process is the state of the optimizer and the process returns empty measurements.

optimizer_fn A tff.learning.optimizers.Optimizer or a no-arg function that returns a tf.keras.optimizers.Optimizer. This optimizer is used to apply client updates to the server model.
model_weights_type A non-federated tff.Type of the model weights to be optimized, which must have a tff.learning.models.ModelWeights container.

A FinalizerProcess that applies the optimizer.

TypeError If value_type does not have a tff.learning.model.sModelWeights Python container, or contains a tff.types.FederatedType.