def test_simple_training(self): it_process = flars_fedavg.build_federated_averaging_process( _keras_model_fn, client_optimizer_fn=lambda: tf.keras.optimizers.SGD(0.1)) server_state = it_process.initialize() # Test out manually setting weights: keras_model = tff.simulation.models.mnist.create_keras_model( compile_model=True) @tf.function def deterministic_batch(): return collections.OrderedDict(x=tf.ones([1, 784], dtype=tf.float32), y=tf.ones([1, 1], dtype=tf.int64)) batch = deterministic_batch() federated_data = [[batch]] def keras_evaluate(state): state.model.assign_weights_to(keras_model) # N.B. The loss computed here won't match the loss computed by TFF because # of the Dropout layer. keras_model.test_on_batch(**batch) loss_list = [] for _ in range(10): keras_evaluate(server_state) server_state, output = it_process.next(server_state, federated_data) loss_list.append(output['loss']) keras_evaluate(server_state) self.assertLess(np.mean(loss_list[7:]), np.mean(loss_list[:2]))
def test_construction(self): it_process = flars_fedavg.build_federated_averaging_process( _keras_model_fn, client_optimizer_fn=lambda: tf.keras.optimizers.SGD(0.1)) self.assertIsInstance(it_process, tff.templates.IterativeProcess) federated_data_type = it_process.next.type_signature.parameter[1] self.assertEqual(str(federated_data_type), '{<x=float32[?,784],y=int64[?,1]>*}@CLIENTS')
def test_self_contained_example_keras_model(self): client_data = create_client_data() train_data = [client_data()] trainer = flars_fedavg.build_federated_averaging_process( _keras_model_fn, client_optimizer_fn=lambda: tf.keras.optimizers.SGD(0.1)) state = trainer.initialize() losses = [] for _ in range(2): state, outputs = trainer.next(state, train_data) # Track the loss. losses.append(outputs['loss']) self.assertLess(losses[1], losses[0])
def _federated_averaging_training_loop(model_fn, client_optimizer_fn, server_optimizer_fn, client_datasets_fn, evaluate_fn, total_rounds=500, rounds_per_eval=1, metrics_hook=None): """A simple example of training loop for the Federated Averaging algorithm. Args: model_fn: A no-arg function that returns a `tff.learning.Model`. client_optimizer_fn: A no-arg function that returns a `tf.keras.optimizers.Optimizer`. server_optimizer_fn: A no-arg function that returns a `tf.keras.optimizers.Optimizer`. client_datasets_fn: A function that takes the round number, and returns a list of `tf.data.Datset`, one per client. evaluate_fn: A function that takes state, performs evaluation, and returns evaluations metrics. total_rounds: Number of rounds to train. rounds_per_eval: How often to call the `metrics_hook` function. metrics_hook: A function taking arguments (training metrics, evaluation metrics, and round number). Optional. Returns: Final `ServerState`. """ logging.info('Starting federated training loop') checkpoint_dir = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name) checkpoint_manager_obj = checkpoint_manager.FileCheckpointManager( checkpoint_dir) if FLAGS.server_optimizer != 'flars': logging.error('Unsupported server_optimzier: %s', FLAGS.server_optimizer) else: iterative_process = flars_fedavg.build_federated_averaging_process( model_fn, client_optimizer_fn=client_optimizer_fn, server_optimizer_fn=server_optimizer_fn) # construct an initial state here to act as a checkpoint template inital_state = iterative_process.initialize() logging.info('Looking for checkpoints in \'%s\'', checkpoint_dir) state, round_num = checkpoint_manager_obj.load_latest_checkpoint(inital_state) if state is None: logging.info('No previous checkpoints, initializing experiment') state = inital_state round_num = 0 if metrics_hook is not None: eval_metrics = evaluate_fn(state) metrics_hook({}, eval_metrics, round_num) checkpoint_manager_obj.save_checkpoint(state, 0) else: logging.info('Restarted from checkpoint round %d', round_num) while round_num < total_rounds: round_num += 1 train_metrics = {} # Reset the executor to clear the cache, and clear the default graph to # garbage collect tf.Functions that will no longer be used. tff.backends.native.set_local_execution_context(max_fanout=25) tf.compat.v1.reset_default_graph() round_start_time = time.time() data_prep_start_time = time.time() train_data = client_datasets_fn(round_num) train_metrics['prepare_datasets_secs'] = time.time() - data_prep_start_time training_start_time = time.time() state, tff_train_metrics = iterative_process.next(state, train_data) tff_train_metrics = tff_train_metrics._asdict(recursive=True) train_metrics.update(tff_train_metrics) train_metrics['training_secs'] = time.time() - training_start_time logging.info('Round {:2d} elapsed time: {:.2f}s .'.format( round_num, (time.time() - round_start_time))) train_metrics['total_round_secs'] = time.time() - round_start_time if (round_num % FLAGS.rounds_per_checkpoint == 0 or round_num == total_rounds): save_checkpoint_start_time = time.time() checkpoint_manager_obj.save_checkpoint(state, round_num) train_metrics['save_checkpoint_secs'] = ( time.time() - save_checkpoint_start_time) if round_num % rounds_per_eval == 0 or round_num == total_rounds: if metrics_hook is not None: eval_metrics = evaluate_fn(state) metrics_hook(train_metrics, eval_metrics, round_num)