def __init__(self, experiment_name, output_dir, hparam_dict): """Returns an initalized `MetricsHook`. Args: experiment_name: A unique filesystem-friendly name for the experiment. output_dir: A root output directory used for all experiment runs in a grid. The `MetricsHook` will combine this with `experiment_name` to form suitable output directories for this run. hparam_dict: A dictionary of hyperparameters to be recorded to .csv and exported to TensorBoard. """ summary_logdir = os.path.join(output_dir, 'logdir/{}'.format(experiment_name)) _check_not_exists(summary_logdir, FLAGS.disable_check_exists) tf.io.gfile.makedirs(summary_logdir) self._summary_writer = tf.summary.create_file_writer( summary_logdir, name=experiment_name) with self._summary_writer.as_default(): hp.hparams(hparam_dict) self._results_file = os.path.join(output_dir, experiment_name, 'results.csv.bz2') # Also write the hparam_dict to a CSV: hparam_dict['results_file'] = self._results_file hparams_file = os.path.join(output_dir, experiment_name, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparams_file) logging.info('Writing ...') logging.info(' result csv to: %s', self._results_file) logging.info(' summaries to: %s', summary_logdir) _check_not_exists(self._results_file, FLAGS.disable_check_exists)
def test_atomic_write_series_with_scalar_data(self, name): series_data = dict(a=1, b=4.0) output_file = os.path.join(absltest.get_default_test_tmpdir(), name) utils_impl.atomic_write_series_to_csv(series_data, output_file) dataframe = pd.read_csv(output_file, index_col=0) pd.testing.assert_frame_equal( pd.DataFrame(pd.Series(series_data), columns=['0']), dataframe)
def _setup_outputs(root_output_dir, experiment_name, hparam_dict): """Set up directories for experiment loops, write hyperparameters to disk.""" if not experiment_name: raise ValueError('experiment_name must be specified.') create_if_not_exists(root_output_dir) checkpoint_dir = os.path.join(root_output_dir, 'checkpoints', experiment_name) create_if_not_exists(checkpoint_dir) checkpoint_mngr = tff.simulation.FileCheckpointManager(checkpoint_dir) results_dir = os.path.join(root_output_dir, 'results', experiment_name) create_if_not_exists(results_dir) csv_file = os.path.join(results_dir, 'experiment.metrics.csv') metrics_mngr = tff.simulation.CSVMetricsManager(csv_file) summary_logdir = os.path.join(root_output_dir, 'logdir', experiment_name) create_if_not_exists(summary_logdir) tensorboard_mngr = tff.simulation.TensorBoardManager(summary_logdir) if hparam_dict: summary_writer = tf.summary.create_file_writer(summary_logdir) hparam_dict['metrics_file'] = metrics_mngr.metrics_filename hparams_file = os.path.join(results_dir, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparams_file) with summary_writer.as_default(): hp.hparams({k: v for k, v in hparam_dict.items() if v is not None}) logging.info('Writing...') logging.info(' checkpoints to: %s', checkpoint_dir) logging.info(' metrics csv to: %s', metrics_mngr.metrics_filename) logging.info(' summaries to: %s', summary_logdir) return checkpoint_mngr, metrics_mngr, tensorboard_mngr
def test_atomic_write_series_with_non_scalar_data(self, name): series_data = dict(a=[1, 2], b=[3.0, 4.0]) output_file = os.path.join(absltest.get_default_test_tmpdir(), name) utils_impl.atomic_write_series_to_csv(series_data, output_file) dataframe = pd.read_csv(output_file, index_col=0) series_data_as_string = dict(a='[1, 2]', b='[3.0, 4.0]') expected_df = pd.DataFrame(pd.Series(series_data_as_string), columns=['0']) pd.testing.assert_frame_equal(expected_df, dataframe)
def _write_hparam_flags(): """Returns an ordered dictionary of pertinent hyperparameter flags.""" hparam_dict = utils_impl.lookup_flag_values(shared_flags) # Update with optimizer flags corresponding to the chosen optimizers. opt_flag_dict = utils_impl.lookup_flag_values(optimizer_flags) opt_flag_dict = optimizer_utils.remove_unused_flags('client', opt_flag_dict) opt_flag_dict = optimizer_utils.remove_unused_flags('server', opt_flag_dict) hparam_dict.update(opt_flag_dict) results_dir = os.path.join(FLAGS.root_output_dir, 'results', FLAGS.experiment_name) utils_impl.create_directory_if_not_exists(results_dir) hparam_file = os.path.join(results_dir, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparam_file)
def _write_hparam_flags(): """Creates an ordered dictionary of hyperparameter flags and writes to CSV.""" hparam_dict = utils_impl.lookup_flag_values(shared_flags) # Update with optimizer flags corresponding to the chosen optimizers. opt_flag_dict = utils_impl.lookup_flag_values(optimizer_flags) opt_flag_dict = optimizer_utils.remove_unused_flags( 'client', opt_flag_dict) opt_flag_dict = optimizer_utils.remove_unused_flags( 'server', opt_flag_dict) if FLAGS.task == 'stackoverflow_nwp_finetune': opt_flag_dict = optimizer_utils.remove_unused_flags( 'finetune', opt_flag_dict) else: opt_flag_dict = optimizer_utils.remove_unused_flags( 'reconstruction', opt_flag_dict) hparam_dict.update(opt_flag_dict) # Update with task-specific flags. task_name = FLAGS.task if task_name in TASK_FLAGS: task_hparam_dict = utils_impl.lookup_flag_values(TASK_FLAGS[task_name]) hparam_dict.update(task_hparam_dict) # Update with finetune flags if FLAGS.task == 'stackoverflow_nwp_finetune': finetune_hparam_dict = utils_impl.lookup_flag_values(finetune_flags) hparam_dict.update(finetune_hparam_dict) # Update with reconstruction flags. recon_hparam_dict = utils_impl.lookup_flag_values(recon_flags) hparam_dict.update(recon_hparam_dict) # Update with DP flags. dp_hparam_dict = utils_impl.lookup_flag_values(dp_flags) hparam_dict.update(dp_hparam_dict) # Update with run flags. run_hparam_dict = utils_impl.lookup_flag_values(run_flags) hparam_dict.update(run_hparam_dict) results_dir = os.path.join(FLAGS.root_output_dir, 'results', FLAGS.experiment_name) utils_impl.create_directory_if_not_exists(results_dir) hparam_file = os.path.join(results_dir, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparam_file)
def run_experiment(): """Data preprocessing and experiment execution.""" emnist_train, _ = emnist_dataset.get_federated_datasets( train_client_batch_size=FLAGS.client_batch_size, train_client_epochs_per_round=FLAGS.client_epochs_per_round, only_digits=False) _, emnist_test = emnist_dataset.get_centralized_datasets() example_dataset = emnist_train.create_tf_dataset_for_client( emnist_train.client_ids[0]) input_spec = example_dataset.element_spec client_datasets_fn = training_utils.build_client_datasets_fn( emnist_train, FLAGS.clients_per_round) evaluate_fn = training_utils.build_centralized_evaluate_fn( eval_dataset=emnist_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) validation_fn = lambda model_weights, round_num: evaluate_fn(model_weights) client_optimizer_fn = functools.partial( utils_impl.create_optimizer_from_flags, 'client') server_optimizer_fn = functools.partial( utils_impl.create_optimizer_from_flags, 'server') def tff_model_fn(): keras_model = model_builder() return tff.learning.from_keras_model(keras_model, input_spec=input_spec, loss=loss_builder(), metrics=metrics_builder()) if FLAGS.use_compression: # We create a `MeasuredProcess` for broadcast process and a # `MeasuredProcess` for aggregate process by providing the # `_broadcast_encoder_fn` and `_mean_encoder_fn` to corresponding utilities. # The fns are called once for each of the model weights created by # tff_model_fn, and return instances of appropriate encoders. encoded_broadcast_process = ( tff.learning.framework.build_encoded_broadcast_process_from_model( tff_model_fn, _broadcast_encoder_fn)) encoded_mean_process = ( tff.learning.framework.build_encoded_mean_process_from_model( tff_model_fn, _mean_encoder_fn)) else: encoded_broadcast_process = None encoded_mean_process = None iterative_process = tff.learning.build_federated_averaging_process( model_fn=tff_model_fn, client_optimizer_fn=client_optimizer_fn, server_optimizer_fn=server_optimizer_fn, aggregation_process=encoded_mean_process, broadcast_process=encoded_broadcast_process) # Log hyperparameters to CSV hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags()) results_dir = os.path.join(FLAGS.root_output_dir, 'results', FLAGS.experiment_name) utils_impl.create_directory_if_not_exists(results_dir) hparam_file = os.path.join(results_dir, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparam_file) training_loop.run(iterative_process=iterative_process, client_datasets_fn=client_datasets_fn, validation_fn=validation_fn, total_rounds=FLAGS.total_rounds, experiment_name=FLAGS.experiment_name, root_output_dir=FLAGS.root_output_dir, rounds_per_eval=FLAGS.rounds_per_eval, rounds_per_checkpoint=FLAGS.rounds_per_checkpoint, rounds_per_profile=FLAGS.rounds_per_profile)
def main(argv): if len(argv) > 1: raise app.UsageError('Expected no command-line arguments, ' 'got: {}'.format(argv)) emnist_train, _ = emnist_dataset.get_federated_datasets( train_client_batch_size=FLAGS.client_batch_size, train_client_epochs_per_round=FLAGS.client_epochs_per_round, only_digits=False) _, emnist_test = emnist_dataset.get_centralized_datasets() if FLAGS.model == 'cnn': model_builder = functools.partial( emnist_models.create_conv_dropout_model, only_digits=False) elif FLAGS.model == '2nn': model_builder = functools.partial( emnist_models.create_two_hidden_layer_model, only_digits=False) else: raise ValueError('Cannot handle model flag [{!s}].'.format(FLAGS.model)) loss_builder = tf.keras.losses.SparseCategoricalCrossentropy metrics_builder = lambda: [tf.keras.metrics.SparseCategoricalAccuracy()] if FLAGS.uniform_weighting: client_weighting = tff.learning.ClientWeighting.UNIFORM else: client_weighting = tff.learning.ClientWeighting.NUM_EXAMPLES def model_fn(): return tff.learning.from_keras_model( model_builder(), loss_builder(), input_spec=emnist_test.element_spec, metrics=metrics_builder()) if FLAGS.noise_multiplier is not None: if not FLAGS.uniform_weighting: raise ValueError( 'Differential privacy is only implemented for uniform weighting.') if FLAGS.noise_multiplier <= 0: raise ValueError('noise_multiplier must be positive if DP is enabled.') if FLAGS.clip is None or FLAGS.clip <= 0: raise ValueError('clip must be positive if DP is enabled.') if not FLAGS.adaptive_clip_learning_rate: aggregation_factory = tff.aggregators.DifferentiallyPrivateFactory.gaussian_fixed( noise_multiplier=FLAGS.noise_multiplier, clients_per_round=FLAGS.clients_per_round, clip=FLAGS.clip) else: if FLAGS.adaptive_clip_learning_rate <= 0: raise ValueError('adaptive_clip_learning_rate must be positive if ' 'adaptive clipping is enabled.') aggregation_factory = tff.aggregators.DifferentiallyPrivateFactory.gaussian_adaptive( noise_multiplier=FLAGS.noise_multiplier, clients_per_round=FLAGS.clients_per_round, initial_l2_norm_clip=FLAGS.clip, target_unclipped_quantile=FLAGS.target_unclipped_quantile, learning_rate=FLAGS.adaptive_clip_learning_rate) else: if FLAGS.uniform_weighting: aggregation_factory = tff.aggregators.UnweightedMeanFactory() else: aggregation_factory = tff.aggregators.MeanFactory() server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('server') client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('client') iterative_process = tff.learning.build_federated_averaging_process( model_fn=model_fn, server_optimizer_fn=server_optimizer_fn, client_weighting=client_weighting, client_optimizer_fn=client_optimizer_fn, model_update_aggregation_factory=aggregation_factory) client_datasets_fn = training_utils.build_client_datasets_fn( emnist_train, FLAGS.clients_per_round) evaluate_fn = training_utils.build_centralized_evaluate_fn( eval_dataset=emnist_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) validation_fn = lambda model_weights, round_num: evaluate_fn(model_weights) logging.info('Training model:') logging.info(model_builder().summary()) # Log hyperparameters to CSV hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags()) results_dir = os.path.join(FLAGS.root_output_dir, 'results', FLAGS.experiment_name) utils_impl.create_directory_if_not_exists(results_dir) hparam_file = os.path.join(results_dir, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparam_file) training_loop.run( iterative_process=iterative_process, client_datasets_fn=client_datasets_fn, validation_fn=validation_fn, total_rounds=FLAGS.total_rounds, experiment_name=FLAGS.experiment_name, root_output_dir=FLAGS.root_output_dir, rounds_per_eval=FLAGS.rounds_per_eval, rounds_per_checkpoint=FLAGS.rounds_per_checkpoint, rounds_per_profile=FLAGS.rounds_per_profile)
def run( keras_model: tf.keras.Model, train_dataset: tf.data.Dataset, experiment_name: str, root_output_dir: str, num_epochs: int, hparams_dict: Optional[Dict[str, Any]] = None, decay_epochs: Optional[int] = None, lr_decay: Optional[float] = None, validation_dataset: Optional[tf.data.Dataset] = None, test_dataset: Optional[tf.data.Dataset] = None ) -> tf.keras.callbacks.History: """Run centralized training for a given compiled `tf.keras.Model`. Args: keras_model: A compiled `tf.keras.Model`. train_dataset: The `tf.data.Dataset` to be used for training. experiment_name: Name of the experiment, used as part of the name of the output directory. root_output_dir: The top-level output directory. The directory `root_output_dir/experiment_name` will contain TensorBoard logs, metrics CSVs and other outputs. num_epochs: How many training epochs to perform. hparams_dict: An optional dict specifying hyperparameters. If provided, the hyperparameters will be written to CSV. decay_epochs: Number of training epochs before decaying the learning rate. lr_decay: How much to decay the learning rate by every `decay_epochs`. validation_dataset: An optional `tf.data.Dataset` used for validation during training. test_dataset: An optional `tf.data.Dataset` used for testing after all training has completed. Returns: A `tf.keras.callbacks.History` object. """ tensorboard_dir = os.path.join(root_output_dir, 'logdir', experiment_name) results_dir = os.path.join(root_output_dir, 'results', experiment_name) for path in [root_output_dir, tensorboard_dir, results_dir]: tf.io.gfile.makedirs(path) if hparams_dict: hparams_file = os.path.join(results_dir, 'hparams.csv') logging.info('Saving hyper parameters to: [%s]', hparams_file) utils_impl.atomic_write_series_to_csv(hparams_dict, hparams_file) csv_logger_callback = keras_callbacks.AtomicCSVLogger(results_dir) tensorboard_callback = tf.keras.callbacks.TensorBoard( log_dir=tensorboard_dir) training_callbacks = [tensorboard_callback, csv_logger_callback] if decay_epochs is not None and decay_epochs > 0: # Reduce the learning rate after a fixed number of epochs. def decay_lr(epoch, learning_rate): if epoch != 0 and epoch % decay_epochs == 0: return learning_rate * lr_decay else: return learning_rate lr_callback = tf.keras.callbacks.LearningRateScheduler(decay_lr, verbose=1) training_callbacks.append(lr_callback) logging.info('Training model:') logging.info(keras_model.summary()) history = keras_model.fit(train_dataset, validation_data=validation_dataset, epochs=num_epochs, callbacks=training_callbacks) logging.info('Final training metrics:') for metric in keras_model.metrics: name = metric.name metric = history.history[name][-1] logging.info('\t%s: %.4f', name, metric) if validation_dataset: logging.info('Final validation metrics:') for metric in keras_model.metrics: name = metric.name metric = history.history['val_{}'.format(name)][-1] logging.info('\t%s: %.4f', name, metric) if test_dataset: test_metrics = keras_model.evaluate(test_dataset, return_dict=True) logging.info('Test metrics:') for metric in keras_model.metrics: name = metric.name metric = test_metrics[name] logging.info('\t%s: %.4f', name, metric) return history
def main(argv): if len(argv) > 1: raise app.UsageError('Expected no command-line arguments, ' 'got: {}'.format(argv)) tff.backends.native.set_local_execution_context(max_fanout=10) model_builder = functools.partial( stackoverflow_models.create_recurrent_model, vocab_size=FLAGS.vocab_size, embedding_size=FLAGS.embedding_size, latent_size=FLAGS.latent_size, num_layers=FLAGS.num_layers, shared_embedding=FLAGS.shared_embedding) loss_builder = functools.partial( tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True) special_tokens = stackoverflow_word_prediction.get_special_tokens( FLAGS.vocab_size) pad_token = special_tokens.pad oov_tokens = special_tokens.oov eos_token = special_tokens.eos def metrics_builder(): return [ keras_metrics.MaskedCategoricalAccuracy(name='accuracy_with_oov', masked_tokens=[pad_token]), keras_metrics.MaskedCategoricalAccuracy(name='accuracy_no_oov', masked_tokens=[pad_token] + oov_tokens), # Notice BOS never appears in ground truth. keras_metrics.MaskedCategoricalAccuracy( name='accuracy_no_oov_or_eos', masked_tokens=[pad_token, eos_token] + oov_tokens), keras_metrics.NumBatchesCounter(), keras_metrics.NumTokensCounter(masked_tokens=[pad_token]), ] train_dataset, _ = stackoverflow_word_prediction.get_federated_datasets( vocab_size=FLAGS.vocab_size, train_client_batch_size=FLAGS.client_batch_size, train_client_epochs_per_round=FLAGS.client_epochs_per_round, max_sequence_length=FLAGS.sequence_length, max_elements_per_train_client=FLAGS.max_elements_per_user) _, validation_dataset, test_dataset = stackoverflow_word_prediction.get_centralized_datasets( vocab_size=FLAGS.vocab_size, max_sequence_length=FLAGS.sequence_length, num_validation_examples=FLAGS.num_validation_examples) if FLAGS.uniform_weighting: client_weighting = tff.learning.ClientWeighting.UNIFORM else: client_weighting = tff.learning.ClientWeighting.NUM_EXAMPLES def model_fn(): return tff.learning.from_keras_model( model_builder(), loss_builder(), input_spec=validation_dataset.element_spec, metrics=metrics_builder()) if FLAGS.noise_multiplier is not None: if not FLAGS.uniform_weighting: raise ValueError( 'Differential privacy is only implemented for uniform weighting.' ) if FLAGS.noise_multiplier <= 0: raise ValueError( 'noise_multiplier must be positive if DP is enabled.') if FLAGS.clip is None or FLAGS.clip <= 0: raise ValueError('clip must be positive if DP is enabled.') if not FLAGS.adaptive_clip_learning_rate: aggregation_factory = tff.aggregators.DifferentiallyPrivateFactory.gaussian_fixed( noise_multiplier=FLAGS.noise_multiplier, clients_per_round=FLAGS.clients_per_round, clip=FLAGS.clip) else: if FLAGS.adaptive_clip_learning_rate <= 0: raise ValueError( 'adaptive_clip_learning_rate must be positive if ' 'adaptive clipping is enabled.') aggregation_factory = tff.aggregators.DifferentiallyPrivateFactory.gaussian_adaptive( noise_multiplier=FLAGS.noise_multiplier, clients_per_round=FLAGS.clients_per_round, initial_l2_norm_clip=FLAGS.clip, target_unclipped_quantile=FLAGS.target_unclipped_quantile, learning_rate=FLAGS.adaptive_clip_learning_rate) else: if FLAGS.uniform_weighting: aggregation_factory = tff.aggregators.UnweightedMeanFactory() else: aggregation_factory = tff.aggregators.MeanFactory() server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags( 'server') client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags( 'client') iterative_process = tff.learning.build_federated_averaging_process( model_fn=model_fn, server_optimizer_fn=server_optimizer_fn, client_weighting=client_weighting, client_optimizer_fn=client_optimizer_fn, model_update_aggregation_factory=aggregation_factory) client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset, FLAGS.clients_per_round) evaluate_fn = training_utils.build_centralized_evaluate_fn( model_builder=model_builder, eval_dataset=validation_dataset, loss_builder=loss_builder, metrics_builder=metrics_builder) validation_fn = lambda state, round_num: evaluate_fn(state.model) evaluate_test_fn = training_utils.build_centralized_evaluate_fn( model_builder=model_builder, # Use both val and test for symmetry with other experiments, which # evaluate on the entire test set. eval_dataset=validation_dataset.concatenate(test_dataset), loss_builder=loss_builder, metrics_builder=metrics_builder) test_fn = lambda state: evaluate_test_fn(state.model) logging.info('Training model:') logging.info(model_builder().summary()) # Log hyperparameters to CSV hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags()) results_dir = os.path.join(FLAGS.root_output_dir, 'results', FLAGS.experiment_name) utils_impl.create_directory_if_not_exists(results_dir) hparam_file = os.path.join(results_dir, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparam_file) training_loop.run(iterative_process=iterative_process, client_datasets_fn=client_datasets_fn, validation_fn=validation_fn, test_fn=test_fn, total_rounds=FLAGS.total_rounds, experiment_name=FLAGS.experiment_name, root_output_dir=FLAGS.root_output_dir, rounds_per_eval=FLAGS.rounds_per_eval, rounds_per_checkpoint=FLAGS.rounds_per_checkpoint)
def run_experiment(): """Data preprocessing and experiment execution.""" emnist_train, _ = emnist_dataset.get_federated_datasets( train_client_batch_size=FLAGS.client_batch_size, train_client_epochs_per_round=FLAGS.client_epochs_per_round, only_digits=FLAGS.only_digits) _, emnist_test = emnist_dataset.get_centralized_datasets( only_digits=FLAGS.only_digits) example_dataset = emnist_train.create_tf_dataset_for_client( emnist_train.client_ids[0]) input_spec = example_dataset.element_spec # Build optimizer functions from flags client_optimizer_fn = functools.partial( utils_impl.create_optimizer_from_flags, 'client') server_optimizer_fn = functools.partial( utils_impl.create_optimizer_from_flags, 'server') def tff_model_fn(): return tff.learning.from_keras_model( keras_model=model_builder(), input_spec=input_spec, loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]) if FLAGS.use_compression: # We create a `tff.templates.MeasuredProcess` for broadcast process and a # `tff.aggregators.WeightedAggregationFactory` for aggregation by providing # the `_broadcast_encoder_fn` and `_mean_encoder_fn` to corresponding # utilities. The fns are called once for each of the model weights created # by tff_model_fn, and return instances of appropriate encoders. encoded_broadcast_process = ( tff.learning.framework.build_encoded_broadcast_process_from_model( tff_model_fn, _broadcast_encoder_fn)) aggregator = tff.aggregators.MeanFactory( tff.aggregators.EncodedSumFactory(_mean_encoder_fn)) else: encoded_broadcast_process = None aggregator = None # Construct the iterative process iterative_process = tff.learning.build_federated_averaging_process( model_fn=tff_model_fn, client_optimizer_fn=client_optimizer_fn, server_optimizer_fn=server_optimizer_fn, broadcast_process=encoded_broadcast_process, model_update_aggregation_factory=aggregator) iterative_process = ( tff.simulation.compose_dataset_computation_with_iterative_process( emnist_train.dataset_computation, iterative_process)) # Create a client sampling function, mapping integer round numbers to lists # of client ids. client_selection_fn = functools.partial( tff.simulation.build_uniform_sampling_fn(emnist_train.client_ids), size=FLAGS.clients_per_round) # Create a validation function evaluate_fn = tff.learning.build_federated_evaluation(tff_model_fn) def validation_fn(state, round_num): if round_num % FLAGS.rounds_per_eval == 0: return evaluate_fn(state.model, [emnist_test]) else: return {} # Log hyperparameters to CSV hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags()) results_dir = os.path.join(FLAGS.root_output_dir, 'results', FLAGS.experiment_name) utils_impl.create_directory_if_not_exists(results_dir) hparam_file = os.path.join(results_dir, 'hparams.csv') utils_impl.atomic_write_series_to_csv(hparam_dict, hparam_file) checkpoint_manager, metrics_managers = _configure_managers() tff.simulation.run_simulation( process=iterative_process, client_selection_fn=client_selection_fn, validation_fn=validation_fn, total_rounds=FLAGS.total_rounds, file_checkpoint_manager=checkpoint_manager, metrics_managers=metrics_managers)