tff.learning.algorithms.build_weighted_fed_avg_with_optimizer_schedule

View source on GitHub

Builds a learning process for FedAvg with client optimizer scheduling.

tff.learning.algorithms.build_weighted_fed_avg_with_optimizer_schedule(
    model_fn: Union[Callable[[], tff.learning.models.VariableModel], tff.learning.models.FunctionalModel],
    client_learning_rate_fn: Callable[[int], float],
    client_optimizer_fn: Callable[[float], tff.learning.optimizers.Optimizer],
    server_optimizer_fn: Optional[tff.learning.optimizers.Optimizer] = None,
    model_distributor: Optional[tff.learning.templates.DistributionProcess] = None,
    model_aggregator: Optional[tff.aggregators.WeightedAggregationFactory] = None,
    metrics_aggregator: Optional[tff.learning.metrics.MetricsAggregatorType] = None,
    loop_implementation: tff.learning.LoopImplementation = tff.learning.LoopImplementation.DATASET_REDUCE
) -> tff.learning.templates.LearningProcess

The primary purpose of this implementation of FedAvg is that it allows for the client optimizer to be scheduled across rounds. Notably, the local optimizaiton step uses a constant learning rate within a round, which is scheduled across rounds of federated training. The process keeps track of how many iterations of .next have occurred (starting at 0), and for each such round_num, the clients will use client_optimizer_fn(round_num) to perform local optimization. This allows learning rate scheduling (eg. starting with a large learning rate and decaying it over time) as well as a small learning rate (eg. switching optimizers as learning progresses).

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

• initialize: A tff.Computation with the functional type signature ( -> S@SERVER), where S is a 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 broadcast to each client using a broadcast function. For each client, local training is performed using client_optimizer_fn. Each client computes the difference between the client model after training and the initial broadcast model. These model deltas are then aggregated at the server using a weighted aggregation function. Clients weighted by the number of examples they see thoughout local training. The aggregate model delta is applied at the server using a server optimizer.