def test_created_keras_optimizer_raises(self): with self.assertRaises(TypeError): model_delta_client_work.build_model_delta_client_work( model_examples.LinearRegression, tf.keras.optimizers.SGD(1.0), client_weighting=client_weight_lib.ClientWeighting.NUM_EXAMPLES )
def test_created_model_raises(self): with self.assertRaises(TypeError): model_delta_client_work.build_model_delta_client_work( model_examples.LinearRegression(), sgdm.build_sgdm(1.0), client_weighting=client_weight_lib.ClientWeighting.NUM_EXAMPLES )
def test_delta_regularizer_yields_smaller_model_delta(self, optimizer): simple_process = model_delta_client_work.build_model_delta_client_work( self.create_model, optimizer, client_weighting=client_weight_lib.ClientWeighting.NUM_EXAMPLES, delta_l2_regularizer=0.0) proximal_process = model_delta_client_work.build_model_delta_client_work( self.create_model, sgdm.build_sgdm(1.0), client_weighting=client_weight_lib.ClientWeighting.NUM_EXAMPLES, delta_l2_regularizer=1.0) client_data = [create_test_dataset()] client_model_weights = [create_test_initial_weights()] simple_output = simple_process.next(simple_process.initialize(), client_model_weights, client_data) proximal_output = proximal_process.next(proximal_process.initialize(), client_model_weights, client_data) simple_update_norm = tf.linalg.global_norm( tf.nest.flatten(simple_output.result[0].update)) proximal_update_norm = tf.linalg.global_norm( tf.nest.flatten(proximal_output.result[0].update)) self.assertGreater(simple_update_norm, proximal_update_norm) self.assertEqual(simple_output.measurements['train']['num_examples'], proximal_output.measurements['train']['num_examples'])
def test_negative_proximal_strength_raises(self): with self.assertRaises(ValueError): model_delta_client_work.build_model_delta_client_work( model_examples.LinearRegression, sgdm.build_sgdm(1.0), client_weighting=client_weight_lib.ClientWeighting. NUM_EXAMPLES, delta_l2_regularizer=-1.0)
def test_custom_metrics_aggregator(self): def sum_then_finalize_then_times_two(metric_finalizers, local_unfinalized_metrics_type): @federated_computation.federated_computation( computation_types.at_clients(local_unfinalized_metrics_type)) def aggregation_computation(client_local_unfinalized_metrics): unfinalized_metrics_sum = intrinsics.federated_sum( client_local_unfinalized_metrics) @tensorflow_computation.tf_computation( local_unfinalized_metrics_type) def finalizer_computation(unfinalized_metrics): finalized_metrics = collections.OrderedDict() for metric_name, metric_finalizer in metric_finalizers.items( ): finalized_metrics[metric_name] = metric_finalizer( unfinalized_metrics[metric_name]) * 2 return finalized_metrics return intrinsics.federated_map(finalizer_computation, unfinalized_metrics_sum) return aggregation_computation process = model_delta_client_work.build_model_delta_client_work( model_fn=self.create_model, optimizer=sgdm.build_sgdm(1.0), client_weighting=client_weight_lib.ClientWeighting.NUM_EXAMPLES, metrics_aggregator=sum_then_finalize_then_times_two) client_model_weights = [create_test_initial_weights()] client_data = [create_test_dataset()] output = process.next(process.initialize(), client_model_weights, client_data) # Train metrics should be multiplied by two by the custom aggregator. self.assertEqual(output.measurements['train']['num_examples'], 16)
def test_execution_with_optimizer(self, optimizer): client_work_process = model_delta_client_work.build_model_delta_client_work( self.create_model, optimizer, client_weighting=client_weight_lib.ClientWeighting.NUM_EXAMPLES) client_data = [create_test_dataset()] client_model_weights = [create_test_initial_weights()] state = client_work_process.initialize() output = client_work_process.next(state, client_model_weights, client_data) self.assertCountEqual(output.measurements.keys(), ['train'])
def test_type_properties(self, optimizer, weighting): model_fn = model_examples.LinearRegression client_work_process = model_delta_client_work.build_model_delta_client_work( model_fn, optimizer, weighting) self.assertIsInstance(client_work_process, client_works.ClientWorkProcess) mw_type = model_utils.ModelWeights( trainable=computation_types.to_type([(tf.float32, (2, 1)), tf.float32]), non_trainable=computation_types.to_type([tf.float32])) expected_param_model_weights_type = computation_types.at_clients( mw_type) expected_param_data_type = computation_types.at_clients( computation_types.SequenceType( computation_types.to_type(model_fn().input_spec))) expected_result_type = computation_types.at_clients( client_works.ClientResult( update=mw_type.trainable, update_weight=computation_types.TensorType(tf.float32))) expected_state_type = computation_types.at_server(()) expected_measurements_type = computation_types.at_server( collections.OrderedDict(train=collections.OrderedDict( loss=tf.float32, num_examples=tf.int32))) expected_initialize_type = computation_types.FunctionType( parameter=None, result=expected_state_type) expected_initialize_type.check_equivalent_to( client_work_process.initialize.type_signature) expected_next_type = computation_types.FunctionType( parameter=collections.OrderedDict( state=expected_state_type, weights=expected_param_model_weights_type, client_data=expected_param_data_type), result=measured_process.MeasuredProcessOutput( expected_state_type, expected_result_type, expected_measurements_type)) expected_next_type.check_equivalent_to( client_work_process.next.type_signature)
def build_basic_fedavg_process(model_fn: Callable[[], model_lib.Model], client_learning_rate: float): """Builds vanilla Federated Averaging process. The created process is the basic form of the Federated Averaging algorithm as proposed by http://proceedings.mlr.press/v54/mcmahan17a/mcmahan17a.pdf in Algorithm 1, for training the model created by `model_fn`. The following is the algorithm in pseudo-code: ``` # Inputs: m: Initial model weights; eta: Client learning rate for i in num_rounds: for c in available_clients_indices: delta_m_c, w_c = client_update(m, eta) aggregate_model_delta = sum_c(model_delta_c * w_c) / sum_c(w_c) m = m - aggregate_model_delta return m # Final trained model. def client_udpate(m, eta): initial_m = m for batch in client_dataset: m = m - eta * grad(m, b) delta_m = initial_m - m return delta_m, size(dataset) ``` The other algorithm hyper parameters (batch size, number of local epochs) are controlled by the data provided to the built process. An example usage of the returned `LearningProcess` in simulation: ``` fedavg = build_basic_fedavg_process(model_fn, 0.1) # Create a `LearningAlgorithmState` containing the initial model weights for # the model returned from `model_fn`. state = fedavg.initialize() for _ in range(num_rounds): client_data = ... # Preprocessed client datasets output = fedavg.next(state, client_data) write_round_metrics(outpus.metrics) # The new state contains the updated model weights after this round. state = output.state ``` Args: model_fn: A no-arg function that returns a `tff.learning.Model`. client_learning_rate: A float. Learning rate for the SGD at clients. Returns: A `LearningProcess`. """ py_typecheck.check_callable(model_fn) py_typecheck.check_type(client_learning_rate, float) @tensorflow_computation.tf_computation() def initial_model_weights_fn(): return model_utils.ModelWeights.from_model(model_fn()) model_weights_type = initial_model_weights_fn.type_signature.result distributor = distributors.build_broadcast_process(model_weights_type) client_work = model_delta_client_work.build_model_delta_client_work( model_fn, sgdm.build_sgdm(client_learning_rate), client_weighting=client_weight_lib.ClientWeighting.NUM_EXAMPLES) aggregator = mean.MeanFactory().create( client_work.next.type_signature.result.result.member.update, client_work.next.type_signature.result.result.member.update_weight) finalizer = finalizers.build_apply_optimizer_finalizer( sgdm.build_sgdm(1.0), model_weights_type) return compose_learning_process(initial_model_weights_fn, distributor, client_work, aggregator, finalizer)
def build_weighted_fed_prox( model_fn: Callable[[], model_lib.Model], proximal_strength: float, client_optimizer_fn: Union[optimizer_base.Optimizer, Callable[[], tf.keras.optimizers.Optimizer]], server_optimizer_fn: Union[optimizer_base.Optimizer, Callable[ [], tf.keras.optimizers.Optimizer]] = DEFAULT_SERVER_OPTIMIZER_FN, client_weighting: Optional[ client_weight_lib.ClientWeighting] = client_weight_lib.ClientWeighting. NUM_EXAMPLES, model_distributor: Optional[distributors.DistributionProcess] = None, model_aggregator: Optional[factory.WeightedAggregationFactory] = None, metrics_aggregator: Optional[Callable[[ model_lib.MetricFinalizersType, computation_types.StructWithPythonType ], computation_base.Computation]] = None, use_experimental_simulation_loop: bool = False ) -> learning_process.LearningProcess: """Builds a learning process that performs the FedProx algorithm. This function creates a `tff.learning.templates.LearningProcess` that performs example-weighted FedProx on client models. This algorithm behaves the same as federated averaging, except that it uses a proximal regularization term that encourages clients to not drift too far from the server model. 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 the 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, as in the FedOpt framework proposed in [Reddi et al., 2021](https://arxiv.org/abs/2003.00295). Note: The default server optimizer function is `tf.keras.optimizers.SGD` with a learning rate of 1.0, which corresponds to adding the model delta to the current server model. This recovers the original FedProx algorithm in [Li et al., 2020](https://arxiv.org/abs/1812.06127). More sophisticated federated averaging procedures may use different learning rates or server optimizers. Args: model_fn: A no-arg function that returns a `tff.learning.Model`. This method 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. proximal_strength: A nonnegative float representing the parameter of FedProx's regularization term. When set to `0`, the algorithm reduces to FedAvg. Higher values prevent clients from moving too far from the server model during local training. client_optimizer_fn: A `tff.learning.optimizers.Optimizer`, or a no-arg callable that returns a `tf.keras.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. 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 broadcasts the model weights on the server to the clients. If set to `None`, the distributor is constructed via `distributors.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.Model.metric_finalizers()`) and a `tff.types.StructWithPythonType` of the unfinalized metrics (i.e., the TFF type of `tff.learning.Model.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: ValueError: If `proximal_parameter` is not a nonnegative float. """ if not isinstance(proximal_strength, float) or proximal_strength < 0.0: raise ValueError( 'proximal_strength must be a nonnegative float, found {}'.format( proximal_strength)) py_typecheck.check_callable(model_fn) @tensorflow_computation.tf_computation() def initial_model_weights_fn(): return model_utils.ModelWeights.from_model(model_fn()) model_weights_type = initial_model_weights_fn.type_signature.result if model_distributor is None: model_distributor = distributors.build_broadcast_process( model_weights_type) if model_aggregator is None: model_aggregator = mean.MeanFactory() py_typecheck.check_type(model_aggregator, factory.WeightedAggregationFactory) aggregator = model_aggregator.create( model_weights_type.trainable, computation_types.TensorType(tf.float32)) process_signature = aggregator.next.type_signature input_client_value_type = process_signature.parameter[1] result_server_value_type = process_signature.result[1] if input_client_value_type.member != result_server_value_type.member: raise TypeError( '`model_update_aggregation_factory` does not produce a ' 'compatible `AggregationProcess`. The processes must ' 'retain the type structure of the inputs on the ' f'server, but got {input_client_value_type.member} != ' f'{result_server_value_type.member}.') if metrics_aggregator is None: metrics_aggregator = metric_aggregator.sum_then_finalize client_work = model_delta_client_work.build_model_delta_client_work( model_fn=model_fn, optimizer=client_optimizer_fn, client_weighting=client_weighting, delta_l2_regularizer=proximal_strength, metrics_aggregator=metrics_aggregator, use_experimental_simulation_loop=use_experimental_simulation_loop) finalizer = finalizers.build_apply_optimizer_finalizer( server_optimizer_fn, model_weights_type) return composers.compose_learning_process(initial_model_weights_fn, model_distributor, client_work, aggregator, finalizer)
def build_weighted_fed_avg( model_fn: Callable[[], model_lib.Model], client_optimizer_fn: Union[optimizer_base.Optimizer, Callable[[], tf.keras.optimizers.Optimizer]], server_optimizer_fn: Union[optimizer_base.Optimizer, Callable[ [], tf.keras.optimizers.Optimizer]] = DEFAULT_SERVER_OPTIMIZER_FN, *, client_weighting: Optional[ client_weight_lib.ClientWeighting] = client_weight_lib.ClientWeighting. NUM_EXAMPLES, model_distributor: Optional[distributors.DistributionProcess] = None, model_aggregator: Optional[factory.WeightedAggregationFactory] = None, metrics_aggregator: Optional[Callable[[ model_lib.MetricFinalizersType, computation_types.StructWithPythonType ], computation_base.Computation]] = None, use_experimental_simulation_loop: bool = False, model: Optional[functional.FunctionalModel] = None, ) -> learning_process.LearningProcess: """Builds a learning process that performs federated averaging. 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. Note: the default server optimizer function is `tf.keras.optimizers.SGD` with a learning rate of 1.0, which corresponds to adding the model delta to the current server model. This recovers the original FedAvg algorithm in [McMahan et al., 2017](https://arxiv.org/abs/1602.05629). More sophisticated federated averaging procedures may use different learning rates or server optimizers (this generalized FedAvg algorithm is described in [Reddi et al., 2021](https://arxiv.org/abs/2003.00295)). Args: model_fn: A no-arg function that returns a `tff.learning.Model`. This method 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. Cannot be used with `model` argument. client_optimizer_fn: A `tff.learning.optimizers.Optimizer`, or a no-arg callable that returns a `tf.keras.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. 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.Model.metric_finalizers()`) and a `tff.types.StructWithPythonType` of the unfinalized metrics (i.e., the TFF type of `tff.learning.Model.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. model: A `tff.learning.models.FunctionalModel` to train. Cannot be used with `model_fn` and must be `None` if `model_fn` is not `None`. Returns: A `tff.learning.templates.LearningProcess`. """ if model is not None and model_fn is not None: raise ValueError( 'Must specify only one of `model` and `model_fn`, both ' 'were not `None`.') elif model is None and model_fn is None: raise ValueError( 'Must specify one of `model` and `model_fn`, both were ' '`None`.') if model_fn is not None: py_typecheck.check_callable(model_fn) else: py_typecheck.check_type(model, functional.FunctionalModel) py_typecheck.check_type(client_weighting, client_weight_lib.ClientWeighting) if model is not None: @tensorflow_computation.tf_computation() def initial_model_weights_fn(): return model_utils.ModelWeights( tuple( tf.convert_to_tensor(w) for w in model.initial_weights[0]), tuple( tf.convert_to_tensor(w) for w in model.initial_weights[1])) else: @tensorflow_computation.tf_computation() def initial_model_weights_fn(): return model_utils.ModelWeights.from_model(model_fn()) model_weights_type = initial_model_weights_fn.type_signature.result if model_distributor is None: model_distributor = distributors.build_broadcast_process( model_weights_type) if model_aggregator is None: model_aggregator = mean.MeanFactory() py_typecheck.check_type(model_aggregator, factory.WeightedAggregationFactory) if model is not None: trainable_weights_type, _ = model_weights_type model_update_type = trainable_weights_type else: model_update_type = model_weights_type.trainable aggregator = model_aggregator.create( model_update_type, computation_types.TensorType(tf.float32)) process_signature = aggregator.next.type_signature input_client_value_type = process_signature.parameter[1] result_server_value_type = process_signature.result[1] if input_client_value_type.member != result_server_value_type.member: raise TypeError( '`model_aggregator` does not produce a compatible ' '`AggregationProcess`. The processes must retain the type ' 'structure of the inputs on the server, but got ' f'{input_client_value_type.member} != ' f'{result_server_value_type.member}.') if metrics_aggregator is None: metrics_aggregator = metric_aggregator.sum_then_finalize if model is not None: client_work = model_delta_client_work.build_functional_model_delta_client_work( model=model, optimizer=client_optimizer_fn, client_weighting=client_weighting, metrics_aggregator=metrics_aggregator) else: client_work = model_delta_client_work.build_model_delta_client_work( model_fn=model_fn, optimizer=client_optimizer_fn, client_weighting=client_weighting, metrics_aggregator=metrics_aggregator, use_experimental_simulation_loop=use_experimental_simulation_loop) finalizer = finalizers.build_apply_optimizer_finalizer( server_optimizer_fn, model_weights_type) return composers.compose_learning_process(initial_model_weights_fn, model_distributor, client_work, aggregator, finalizer)