Learn the latest in machine learning, generative AI, and more at WiML Symposium 2023Register

tff.learning.algorithms.build_weighted_fed_avg

Builds a learning process that performs federated averaging.

tff.learning.algorithms.build_weighted_fed_avg(
    model_fn: Union[Callable[[], tff.learning.models.VariableModel], tff.learning.models.FunctionalModel],
    client_optimizer_fn: Union[tff.learning.optimizers.Optimizer, Callable[[], tf.keras.optimizers.Optimizer]],
    server_optimizer_fn: Union[tff.learning.optimizers.Optimizer, Callable[[], tf.keras.optimizers.Optimizer]] = DEFAULT_SERVER_OPTIMIZER_FN,
    *,
    client_weighting: Optional[tff.learning.ClientWeighting] = tff.learning.ClientWeighting.NUM_EXAMPLES,
    model_distributor: Optional[tff.learning.templates.DistributionProcess] = None,
    model_aggregator: Optional[tff.aggregators.WeightedAggregationFactory] = None,
    metrics_aggregator: Optional[tff.learning.metrics.MetricsAggregatorType] = None,
    use_experimental_simulation_loop: bool = False
) -> tff.learning.templates.LearningProcess

Used in the notebooks

Used in the tutorials
Federated Learning for Image Classification TFF for Federated Learning Research: Model and Update Compression Federated Learning for Text Generation High-performance simulations with TFF TFF simulations with accelerators

This function creates a tff.learning.templates.LearningProcess that performs federated averaging on client models. The iterative process has the following methods inherited from tff.learning.templates.LearningProcess:

initialize: A tff.Computation with the functional type signature ( -> S@SERVER), where S is a tff.learning.templates.LearningAlgorithmState representing the initial state of the server.
next: A tff.Computation with the functional type signature (<S@SERVER, {B*}@CLIENTS> -> <L@SERVER>) where S is a tff.learning.templates.LearningAlgorithmState whose type matches the output of initialize and {B*}@CLIENTS represents the client datasets. The output L contains the updated server state, as well as aggregated metrics at the server, including client training metrics and any other metrics from distribution and aggregation processes.
get_model_weights: A tff.Computation with type signature (S -> M), where S is a tff.learning.templates.LearningAlgorithmState whose type matches the output of initialize and next, and M represents the type of the model weights used during training.
set_model_weights: A tff.Computation with type signature (<S, M> -> S), where S is a tff.learning.templates.LearningAlgorithmState whose type matches the output of initialize and M represents the type of the model weights used during training.

Each time the next method is called, the server model is communicated to each client using the provided model_distributor. For each client, local training is performed using client_optimizer_fn. Each client computes the difference between the client model after training and its initial model. These model deltas are then aggregated at the server using a weighted aggregation function, according to client_weighting. The aggregate model delta is applied at the server using a server optimizer.

Args
`model_fn`	A no-arg function that returns a `tff.learning.models.VariableModel`, or an instance of a `tff.learning.models.FunctionalModel`. When passing a callable, the callable must not capture TensorFlow tensors or variables and use them. The model must be constructed entirely from scratch on each invocation, returning the same pre-constructed model each call will result in an error.
`client_optimizer_fn`	A `tff.learning.optimizers.Optimizer`, or a no-arg callable that returns a `tf.keras.Optimizer`. If `model_fn` is a `tff.learning.models.FunctionalModel`, must be a `tff.learning.optimizers.Optimizer`.
`server_optimizer_fn`	A `tff.learning.optimizers.Optimizer`, or a no-arg callable that returns a `tf.keras.Optimizer`. By default, this uses `tf.keras.optimizers.SGD` with a learning rate of 1.0. If `model_fn` is a `tff.learning.models.FunctionalModel`, must be a `tff.learning.optimizers.Optimizer`.
`client_weighting`	A member of `tff.learning.ClientWeighting` that specifies a built-in weighting method. By default, weighting by number of examples is used.
`model_distributor`	An optional `DistributionProcess` that distributes the model weights on the server to the clients. If set to `None`, the distributor is constructed via `tff.learning.templates.build_broadcast_process`.
`model_aggregator`	An optional `tff.aggregators.WeightedAggregationFactory` used to aggregate client updates on the server. If `None`, this is set to `tff.aggregators.MeanFactory`.
`metrics_aggregator`	A function that takes in the metric finalizers (i.e., `tff.learning.models.VariableModel.metric_finalizers()`) and a `tff.types.StructWithPythonType` of the unfinalized metrics (i.e., the TFF type of `tff.learning.models.VariableModel.report_local_unfinalized_metrics()`), and returns a `tff.Computation` for aggregating the unfinalized metrics. If `None`, this is set to `tff.learning.metrics.sum_then_finalize`.
`use_experimental_simulation_loop`	Controls the reduce loop function for input dataset. An experimental reduce loop is used for simulation. It is currently necessary to set this flag to True for performant GPU simulations.

Returns
A `tff.learning.templates.LearningProcess`.

Raises
`TypeError`	If arguments are not the documented types.

tff.learning.algorithms.build_weighted_fed_avg

Used in the notebooks

Args

Returns

Raises