def test_evaluate_fn_with_tuple_of_trainable_variables(self): loss_builder = tf.keras.losses.MeanSquaredError metrics_builder = lambda: [tf.keras.metrics.MeanSquaredError()] def tff_model_fn(): return tff.learning.from_keras_model( keras_model=model_builder(), input_spec=get_input_spec(), loss=loss_builder(), metrics=metrics_builder()) iterative_process = tff.learning.build_federated_averaging_process( tff_model_fn, client_optimizer_fn=tf.keras.optimizers.SGD) state = iterative_process.initialize() test_dataset = create_tf_dataset_for_client(1) reference_model = tff.learning.ModelWeights( trainable=tuple(state.model.trainable), non_trainable=tuple(state.model.non_trainable)) evaluate_fn = training_utils.build_evaluate_fn(test_dataset, model_builder, loss_builder, metrics_builder) test_metrics = evaluate_fn(reference_model) self.assertIn('loss', test_metrics)
def run_federated( iterative_process_builder: Callable[..., tff.templates.IterativeProcess], client_epochs_per_round: int, client_batch_size: int, clients_per_round: int, schedule: Optional[str] = 'none', beta: Optional[float] = 0., max_batches_per_client: Optional[int] = -1, client_datasets_random_seed: Optional[int] = None, crop_size: Optional[int] = 24, total_rounds: Optional[int] = 1500, experiment_name: Optional[str] = 'federated_cifar100', root_output_dir: Optional[str] = '/tmp/fed_opt', max_eval_batches: Optional[int] = None, **kwargs): """Runs an iterative process on the CIFAR-100 classification task. This method will load and pre-process dataset and construct a model used for the task. It then uses `iterative_process_builder` to create an iterative process that it applies to the task, using `federated_research.utils.training_loop`. We assume that the iterative process has the following functional type signatures: * `initialize`: `( -> S@SERVER)` where `S` represents the server state. * `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S` represents the server state, `{B*}` represents the client datasets, and `T` represents a python `Mapping` object. Moreover, the server state must have an attribute `model` of type `tff.learning.ModelWeights`. Args: iterative_process_builder: A function that accepts a no-arg `model_fn`, and returns a `tff.templates.IterativeProcess`. The `model_fn` must return a `tff.learning.Model`. client_epochs_per_round: An integer representing the number of epochs of training performed per client in each training round. client_batch_size: An integer representing the batch size used on clients. clients_per_round: An integer representing the number of clients participating in each round. max_batches_per_client: An optional int specifying the number of batches taken by each client at each round. If `-1`, the entire client dataset is used. client_datasets_random_seed: An optional int used to seed which clients are sampled at each round. If `None`, no seed is used. crop_size: An optional integer representing the resulting size of input images after preprocessing. total_rounds: The number of federated training rounds. experiment_name: The name of the experiment being run. This will be appended to the `root_output_dir` for purposes of writing outputs. root_output_dir: The name of the root output directory for writing experiment outputs. max_eval_batches: If set to a positive integer, evaluation datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full evaluation datasets are used. **kwargs: Additional arguments configuring the training loop. For details on supported arguments, see `federated_research/utils/training_utils.py`. """ crop_shape = (crop_size, crop_size, 3) cifar_train, _, fed_test_data = cifar100_dataset.get_federated_cifar100( client_epochs_per_round=client_epochs_per_round, train_batch_size=client_batch_size, crop_shape=crop_shape, max_batches_per_client=max_batches_per_client) _, cifar_test = cifar100_dataset.get_centralized_datasets( train_batch_size=client_batch_size, max_test_batches=max_eval_batches, crop_shape=crop_shape) input_spec = cifar_train.create_tf_dataset_for_client( cifar_train.client_ids[0]).element_spec model_builder = functools.partial(resnet_models.create_resnet18, input_shape=crop_shape, num_classes=NUM_CLASSES) loss_builder = tf.keras.losses.SparseCategoricalCrossentropy metrics_builder = lambda: [tf.keras.metrics.SparseCategoricalAccuracy()] def tff_model_fn() -> tff.learning.Model: return tff.learning.from_keras_model(keras_model=model_builder(), input_spec=input_spec, loss=loss_builder(), metrics=metrics_builder()) training_process = iterative_process_builder(tff_model_fn) evaluate_fn = training_utils.build_evaluate_fn( eval_dataset=cifar_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) test_fn = training_utils.build_unweighted_test_fn( federated_eval_dataset=fed_test_data, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) logging.info('Training model:') logging.info(model_builder().summary()) try: var = kwargs['hparam_dict']['var_q_clients'] q_client = np.load( f'/home/monica/AVAIL_VECTORS/q_client_{var}_cifar.npy') except: logging.info( 'Could not load q_client - initializing random availabilities') q_client = None if schedule == 'none': client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=cifar_train, train_clients_per_round=clients_per_round, random_seed=client_datasets_random_seed, var_q_clients=kwargs['hparam_dict']['var_q_clients'], f_mult=kwargs['hparam_dict']['f_mult'], f_intercept=kwargs['hparam_dict']['f_intercept'], sine_wave=kwargs['hparam_dict']['sine_wave'], use_p=True, q_client=q_client, ) training_loop.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs) elif schedule == 'loss': if 'loss_pool_size' in kwargs['hparam_dict'] and kwargs['hparam_dict'][ 'loss_pool_size'] is not None: loss_pool_size = kwargs['hparam_dict']['loss_pool_size'] logging.info(f'Loss pool size: {loss_pool_size}') client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=cifar_train, train_clients_per_round=loss_pool_size, random_seed=client_datasets_random_seed, var_q_clients=kwargs['hparam_dict']['var_q_clients'], f_mult=kwargs['hparam_dict']['f_mult'], f_intercept=kwargs['hparam_dict']['f_intercept'], sine_wave=kwargs['hparam_dict']['sine_wave'], use_p=True, q_client=q_client, ) training_loop_loss.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, total_clients=loss_pool_size, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs) else: raise ValueError('Loss pool size not specified') else: client_datasets_fn = training_utils.build_availability_client_datasets_fn( train_dataset=cifar_train, train_clients_per_round=clients_per_round, random_seed=client_datasets_random_seed, beta=beta, var_q_clients=kwargs['hparam_dict']['var_q_clients'], f_mult=kwargs['hparam_dict']['f_mult'], f_intercept=kwargs['hparam_dict']['f_intercept'], sine_wave=kwargs['hparam_dict']['sine_wave'], q_client=q_client, ) training_loop_importance.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs)
def run_federated( iterative_process_builder: Callable[..., tff.templates.IterativeProcess], client_epochs_per_round: int, client_batch_size: int, clients_per_round: int, max_batches_per_client: Optional[int] = -1, client_datasets_random_seed: Optional[int] = None, total_rounds: Optional[int] = 1500, experiment_name: Optional[str] = 'federated_emnist_ae', root_output_dir: Optional[str] = '/tmp/fed_opt', max_eval_batches: Optional[int] = None, **kwargs): """Runs an iterative process on the EMNIST autoencoder task. This method will load and pre-process dataset and construct a model used for the task. It then uses `iterative_process_builder` to create an iterative process that it applies to the task, using `federated_research.utils.training_loop`. We assume that the iterative process has the following functional type signatures: * `initialize`: `( -> S@SERVER)` where `S` represents the server state. * `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S` represents the server state, `{B*}` represents the client datasets, and `T` represents a python `Mapping` object. Moreover, the server state must have an attribute `model` of type `tff.learning.ModelWeights`. Args: iterative_process_builder: A function that accepts a no-arg `model_fn`, and returns a `tff.templates.IterativeProcess`. The `model_fn` must return a `tff.learning.Model`. client_epochs_per_round: An integer representing the number of epochs of training performed per client in each training round. client_batch_size: An integer representing the batch size used on clients. clients_per_round: An integer representing the number of clients participating in each round. max_batches_per_client: An optional int specifying the number of batches taken by each client at each round. If `-1`, the entire client dataset is used. client_datasets_random_seed: An optional int used to seed which clients are sampled at each round. If `None`, no seed is used. total_rounds: The number of federated training rounds. experiment_name: The name of the experiment being run. This will be appended to the `root_output_dir` for purposes of writing outputs. root_output_dir: The name of the root output directory for writing experiment outputs. max_eval_batches: If set to a positive integer, evaluation datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full evaluation datasets are used. **kwargs: Additional arguments configuring the training loop. For details on supported arguments, see `federated_research/utils/training_utils.py`. """ emnist_train, _ = emnist_ae_dataset.get_emnist_datasets( client_batch_size=client_batch_size, client_epochs_per_round=client_epochs_per_round, max_batches_per_client=max_batches_per_client, only_digits=False) _, emnist_test = emnist_ae_dataset.get_centralized_datasets( train_batch_size=client_batch_size, max_test_batches=max_eval_batches, only_digits=False) input_spec = emnist_train.create_tf_dataset_for_client( emnist_train.client_ids[0]).element_spec model_builder = emnist_ae_models.create_autoencoder_model loss_builder = functools.partial(tf.keras.losses.MeanSquaredError, reduction=tf.keras.losses.Reduction.SUM) metrics_builder = lambda: [tf.keras.metrics.MeanSquaredError()] def tff_model_fn() -> tff.learning.Model: return tff.learning.from_keras_model(keras_model=model_builder(), input_spec=input_spec, loss=loss_builder(), metrics=metrics_builder()) training_process = iterative_process_builder(tff_model_fn) client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=emnist_train, train_clients_per_round=clients_per_round, random_seed=client_datasets_random_seed) evaluate_fn = training_utils.build_evaluate_fn( eval_dataset=emnist_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) logging.info('Training model:') logging.info(model_builder().summary()) training_loop.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=evaluate_fn, total_rounds=total_rounds, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs)
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_dataset.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]), ] datasets = stackoverflow_dataset.construct_word_level_datasets( FLAGS.vocab_size, FLAGS.client_batch_size, FLAGS.client_epochs_per_round, FLAGS.sequence_length, FLAGS.max_elements_per_user, FLAGS.num_validation_examples) train_dataset, validation_dataset, test_dataset = datasets if FLAGS.uniform_weighting: def client_weight_fn(local_outputs): del local_outputs return 1.0 else: def client_weight_fn(local_outputs): return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32) 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.' ) dp_query = tff.utils.build_dp_query( clip=FLAGS.clip, noise_multiplier=FLAGS.noise_multiplier, expected_total_weight=FLAGS.clients_per_round, adaptive_clip_learning_rate=FLAGS.adaptive_clip_learning_rate, target_unclipped_quantile=FLAGS.target_unclipped_quantile, clipped_count_budget_allocation=FLAGS. clipped_count_budget_allocation, expected_clients_per_round=FLAGS.clients_per_round, per_vector_clipping=FLAGS.per_vector_clipping, model=model_fn()) weights_type = tff.learning.framework.weights_type_from_model(model_fn) aggregation_process = tff.utils.build_dp_aggregate_process( weights_type.trainable, dp_query) else: aggregation_process = None 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_weight_fn=client_weight_fn, client_optimizer_fn=client_optimizer_fn, aggregation_process=aggregation_process) client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset, FLAGS.clients_per_round) evaluate_fn = training_utils.build_evaluate_fn( model_builder=model_builder, eval_dataset=validation_dataset, loss_builder=loss_builder, metrics_builder=metrics_builder) test_fn = training_utils.build_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) logging.info('Training model:') logging.info(model_builder().summary()) hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags()) training_loop_dict = utils_impl.lookup_flag_values(training_loop_flags) training_loop.run(iterative_process=iterative_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, hparam_dict=hparam_dict, **training_loop_dict)
def run_federated( iterative_process_builder: Callable[..., tff.templates.IterativeProcess], client_epochs_per_round: int, client_batch_size: int, clients_per_round: int, max_batches_per_client: Optional[int] = -1, client_datasets_random_seed: Optional[int] = None, sequence_length: Optional[int] = 80, total_rounds: Optional[int] = 1500, experiment_name: Optional[str] = 'federated_shakespeare', root_output_dir: Optional[str] = '/tmp/fed_opt', max_eval_batches: Optional[int] = None, **kwargs): """Runs an iterative process on a Shakespeare next character prediction task. This method will load and pre-process dataset and construct a model used for the task. It then uses `iterative_process_builder` to create an iterative process that it applies to the task, using `federated_research.utils.training_loop`. We assume that the iterative process has the following functional type signatures: * `initialize`: `( -> S@SERVER)` where `S` represents the server state. * `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S` represents the server state, `{B*}` represents the client datasets, and `T` represents a python `Mapping` object. Moreover, the server state must have an attribute `model` of type `tff.learning.ModelWeights`. Args: iterative_process_builder: A function that accepts a no-arg `model_fn`, and a `client_weight_fn`, and returns a `tff.templates.IterativeProcess`. The `model_fn` must return a `tff.learning.Model`. client_epochs_per_round: An integer representing the number of epochs of training performed per client in each training round. client_batch_size: An integer representing the batch size used on clients. clients_per_round: An integer representing the number of clients participating in each round. max_batches_per_client: An optional int specifying the number of batches taken by each client at each round. If `-1`, the entire client dataset is used. client_datasets_random_seed: An optional int used to seed which clients are sampled at each round. If `None`, no seed is used. sequence_length: An int specifying the length of the character sequences used for prediction. total_rounds: The number of federated training rounds. experiment_name: The name of the experiment being run. This will be appended to the `root_output_dir` for purposes of writing outputs. root_output_dir: The name of the root output directory for writing experiment outputs. max_eval_batches: If set to a positive integer, evaluation datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full evaluation datasets are used. **kwargs: Additional arguments configuring the training loop. For details on supported arguments, see `federated_research/utils/training_utils.py`. """ train_clientdata = shakespeare_dataset.construct_character_level_datasets( client_batch_size=client_batch_size, client_epochs_per_round=client_epochs_per_round, sequence_length=sequence_length, max_batches_per_client=max_batches_per_client) _, test_dataset = shakespeare_dataset.get_centralized_datasets( train_batch_size=client_batch_size, max_test_batches=max_eval_batches, sequence_length=sequence_length) model_builder = functools.partial(create_shakespeare_model, sequence_length=sequence_length) loss_builder = functools.partial( tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True) input_spec = train_clientdata.create_tf_dataset_for_client( train_clientdata.client_ids[0]).element_spec def client_weight_fn(local_outputs): # Num_tokens is a tensor with type int64[1], to use as a weight need # a float32 scalar. return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32) def tff_model_fn() -> tff.learning.Model: return tff.learning.from_keras_model(keras_model=model_builder(), input_spec=input_spec, loss=loss_builder(), metrics=metrics_builder()) training_process = iterative_process_builder( tff_model_fn, client_weight_fn=client_weight_fn) client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=train_clientdata, train_clients_per_round=clients_per_round, random_seed=client_datasets_random_seed) evaluate_fn = training_utils.build_evaluate_fn( eval_dataset=test_dataset, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) logging.info('Training model:') logging.info(model_builder().summary()) training_loop.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=evaluate_fn, total_rounds=total_rounds, total_clients=clients_per_round, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs)
def main(argv): if len(argv) > 1: raise app.UsageError('Expected no command-line arguments, ' 'got: {}'.format(argv)) emnist_train, emnist_test = emnist_dataset.get_emnist_datasets( FLAGS.client_batch_size, FLAGS.client_epochs_per_round, only_digits=False) 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: def client_weight_fn(local_outputs): del local_outputs return 1.0 else: client_weight_fn = None # Defaults to the number of examples per client. 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.' ) dp_query = tff.utils.build_dp_query( clip=FLAGS.clip, noise_multiplier=FLAGS.noise_multiplier, expected_total_weight=FLAGS.clients_per_round, adaptive_clip_learning_rate=FLAGS.adaptive_clip_learning_rate, target_unclipped_quantile=FLAGS.target_unclipped_quantile, clipped_count_budget_allocation=FLAGS. clipped_count_budget_allocation, expected_clients_per_round=FLAGS.clients_per_round, per_vector_clipping=FLAGS.per_vector_clipping, model=model_fn()) weights_type = tff.learning.framework.weights_type_from_model(model_fn) aggregation_process = tff.utils.build_dp_aggregate_process( weights_type.trainable, dp_query) else: aggregation_process = None 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_weight_fn=client_weight_fn, client_optimizer_fn=client_optimizer_fn, aggregation_process=aggregation_process) client_datasets_fn = training_utils.build_client_datasets_fn( emnist_train, FLAGS.clients_per_round) evaluate_fn = training_utils.build_evaluate_fn( eval_dataset=emnist_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) logging.info('Training model:') logging.info(model_builder().summary()) hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags()) training_loop_dict = utils_impl.lookup_flag_values(training_loop_flags) training_loop.run(iterative_process=iterative_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, hparam_dict=hparam_dict, **training_loop_dict)
def run_federated( iterative_process_builder: Callable[..., tff.templates.IterativeProcess], client_epochs_per_round: int, client_batch_size: int, clients_per_round: int, schedule: Optional[str] = 'none', beta: Optional[float] = 0., max_batches_per_client: Optional[int] = -1, client_datasets_random_seed: Optional[int] = None, model: Optional[str] = 'cnn', total_rounds: Optional[int] = 1500, experiment_name: Optional[str] = 'federated_synthetic', root_output_dir: Optional[str] = '/tmp/fed_opt', max_eval_batches: Optional[int] = None, alpha: Optional[float] = 0., beta_data: Optional[float] = 0., iid: Optional[int] = 0, num_users: Optional[int] = 100, **kwargs): """Runs an iterative process on the EMNIST character recognition task. This method will load and pre-process dataset and construct a model used for the task. It then uses `iterative_process_builder` to create an iterative process that it applies to the task, using `federated_research.utils.training_loop`. We assume that the iterative process has the following functional type signatures: * `initialize`: `( -> S@SERVER)` where `S` represents the server state. * `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S` represents the server state, `{B*}` represents the client datasets, and `T` represents a python `Mapping` object. Moreover, the server state must have an attribute `model` of type `tff.learning.ModelWeights`. Args: iterative_process_builder: A function that accepts a no-arg `model_fn`, and returns a `tff.templates.IterativeProcess`. The `model_fn` must return a `tff.learning.Model`. client_epochs_per_round: An integer representing the number of epochs of training performed per client in each training round. client_batch_size: An integer representing the batch size used on clients. clients_per_round: An integer representing the number of clients participating in each round. max_batches_per_client: An optional int specifying the number of batches taken by each client at each round. If `-1`, the entire client dataset is used. client_datasets_random_seed: An optional int used to seed which clients are sampled at each round. If `None`, no seed is used. model: A string specifying the model used for character recognition. Can be one of `cnn` and `2nn`, corresponding to a CNN model and a densely connected 2-layer model (respectively). total_rounds: The number of federated training rounds. experiment_name: The name of the experiment being run. This will be appended to the `root_output_dir` for purposes of writing outputs. root_output_dir: The name of the root output directory for writing experiment outputs. max_eval_batches: If set to a positive integer, evaluation datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full evaluation datasets are used. **kwargs: Additional arguments configuring the training loop. For details on supported arguments, see `federated_research/utils/training_utils.py`. """ logging.info(f' DATA PARAMS: ') logging.info(f' Num Users: {num_users}') logging.info(f' alpha: {alpha}') logging.info(f' beta: {beta_data}') logging.info(f' iid: {iid}') train_data, test_data, federated_test = synthetic_dataset.generate_federated_softmax_data( batch_size=client_batch_size, client_epochs_per_round=client_epochs_per_round, test_batch_size=100, alpha=alpha, beta=beta_data, iid=iid, num_users=num_users) input_spec = train_data.create_tf_dataset_for_client( train_data.client_ids[0]).element_spec model_builder = functools.partial(create_logistic_regression_model) loss_builder = tf.keras.losses.SparseCategoricalCrossentropy metrics_builder = lambda: [tf.keras.metrics.SparseCategoricalAccuracy()] def tff_model_fn() -> tff.learning.Model: return tff.learning.from_keras_model(keras_model=model_builder(), input_spec=input_spec, loss=loss_builder(), metrics=metrics_builder()) training_process = iterative_process_builder(tff_model_fn) evaluate_fn = training_utils.build_evaluate_fn( eval_dataset=test_data, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) test_fn = training_utils.build_unweighted_test_fn( federated_eval_dataset=federated_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) logging.info('Training model:') logging.info(model_builder().summary()) try: var = kwargs['hparam_dict']['var_q_clients'] print(f'Variance: {var}') q_client = np.load( f'/home/monica/AVAIL_VECTORS/q_client_{var}_synthetic.npy') #if q_client is None: #logging.info('Could not load q_client - initializing random availabilities') #q_client=None except: logging.info( 'Could not load q_client - initializing random availabilities') q_client = None if schedule == 'none': client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=train_data, train_clients_per_round=clients_per_round, random_seed=client_datasets_random_seed, min_clients=kwargs['hparam_dict']['min_clients'], var_q_clients=kwargs['hparam_dict']['var_q_clients'], f_mult=kwargs['hparam_dict']['f_mult'], f_intercept=kwargs['hparam_dict']['f_intercept'], sine_wave=kwargs['hparam_dict']['sine_wave'], use_p=True, q_client=q_client, ) training_loop.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs) elif schedule == 'loss': if 'loss_pool_size' in kwargs['hparam_dict'] and kwargs['hparam_dict'][ 'loss_pool_size'] is not None: loss_pool_size = kwargs['hparam_dict']['loss_pool_size'] logging.info(f'Loss pool size: {loss_pool_size}') client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=train_data, train_clients_per_round=loss_pool_size, random_seed=client_datasets_random_seed, min_clients=kwargs['hparam_dict']['min_clients'], var_q_clients=kwargs['hparam_dict']['var_q_clients'], f_mult=kwargs['hparam_dict']['f_mult'], f_intercept=kwargs['hparam_dict']['f_intercept'], sine_wave=kwargs['hparam_dict']['sine_wave'], use_p=True, q_client=q_client) training_loop_loss.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, total_clients=loss_pool_size, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs) else: raise ValueError('Loss pool size not specified') else: init_p = kwargs['hparam_dict']['initialize_p'] logging.info(f'Initializing as p = {init_p}') client_datasets_fn = training_utils.build_availability_client_datasets_fn( train_dataset=train_data, train_clients_per_round=clients_per_round, beta=beta, min_clients=kwargs['hparam_dict']['min_clients'], var_q_clients=kwargs['hparam_dict']['var_q_clients'], f_mult=kwargs['hparam_dict']['f_mult'], f_intercept=kwargs['hparam_dict']['f_intercept'], sine_wave=kwargs['hparam_dict']['sine_wave'], q_client=q_client, initialize_p=init_p, ) training_loop_importance.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs)
def run_federated( iterative_process_builder: Callable[..., tff.templates.IterativeProcess], client_epochs_per_round: int, client_batch_size: int, clients_per_round: int, max_batches_per_client: Optional[int] = -1, client_datasets_random_seed: Optional[int] = None, vocab_size: Optional[int] = 10000, num_oov_buckets: Optional[int] = 1, sequence_length: Optional[int] = 20, max_elements_per_user: Optional[int] = 1000, num_validation_examples: Optional[int] = 10000, embedding_size: Optional[int] = 96, latent_size: Optional[int] = 670, num_layers: Optional[int] = 1, shared_embedding: Optional[bool] = False, total_rounds: Optional[int] = 1500, experiment_name: Optional[str] = 'federated_so_nwp', root_output_dir: Optional[str] = '/tmp/fed_opt', max_eval_batches: Optional[int] = None, **kwargs): """Runs an iterative process on the Stack Overflow next word prediction task. This method will load and pre-process dataset and construct a model used for the task. It then uses `iterative_process_builder` to create an iterative process that it applies to the task, using `federated_research.utils.training_loop`. We assume that the iterative process has the following functional type signatures: * `initialize`: `( -> S@SERVER)` where `S` represents the server state. * `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S` represents the server state, `{B*}` represents the client datasets, and `T` represents a python `Mapping` object. Moreover, the server state must have an attribute `model` of type `tff.learning.ModelWeights`. Args: iterative_process_builder: A function that accepts a no-arg `model_fn`, a `client_weight_fn` and returns a `tff.templates.IterativeProcess`. The `model_fn` must return a `tff.learning.Model`. client_epochs_per_round: An integer representing the number of epochs of training performed per client in each training round. client_batch_size: An integer representing the batch size used on clients. clients_per_round: An integer representing the number of clients participating in each round. max_batches_per_client: An optional int specifying the number of batches taken by each client at each round. If `-1`, the entire client dataset is used. client_datasets_random_seed: An optional int used to seed which clients are sampled at each round. If `None`, no seed is used. vocab_size: Integer dictating the number of most frequent words to use in the vocabulary. num_oov_buckets: The number of out-of-vocabulary buckets to use. sequence_length: The maximum number of words to take for each sequence. max_elements_per_user: The maximum number of elements processed for each client's dataset. num_validation_examples: The number of test examples to use for validation. embedding_size: The dimension of the word embedding layer. latent_size: The dimension of the latent units in the recurrent layers. num_layers: The number of stacked recurrent layers to use. shared_embedding: Boolean indicating whether to tie input and output embeddings. total_rounds: The number of federated training rounds. experiment_name: The name of the experiment being run. This will be appended to the `root_output_dir` for purposes of writing outputs. root_output_dir: The name of the root output directory for writing experiment outputs. max_eval_batches: If set to a positive integer, evaluation datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full evaluation datasets are used. **kwargs: Additional arguments configuring the training loop. For details on supported arguments, see `federated_research/utils/training_utils.py`. """ model_builder = functools.partial( stackoverflow_models.create_recurrent_model, vocab_size=vocab_size, num_oov_buckets=num_oov_buckets, embedding_size=embedding_size, latent_size=latent_size, num_layers=num_layers, shared_embedding=shared_embedding) loss_builder = functools.partial( tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True) special_tokens = stackoverflow_dataset.get_special_tokens( vocab_size, num_oov_buckets) 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_clientdata, _, _ = tff.simulation.datasets.stackoverflow.load_data() # TODO(b/161914546): consider moving evaluation to use # `tff.learning.build_federated_evaluation` to get metrics over client # distributions, as well as the example weight means from this centralized # evaluation. _, validation_dataset, test_dataset = stackoverflow_dataset.get_centralized_datasets( vocab_size=vocab_size, max_seq_len=sequence_length, train_batch_size=client_batch_size, max_validation_batches=max_eval_batches, max_test_batches=max_eval_batches, num_validation_examples=num_validation_examples, num_oov_buckets=num_oov_buckets) train_dataset_preprocess_comp = stackoverflow_dataset.create_train_dataset_preprocess_fn( vocab=stackoverflow_dataset.create_vocab(vocab_size), num_oov_buckets=num_oov_buckets, client_batch_size=client_batch_size, client_epochs_per_round=client_epochs_per_round, max_seq_len=sequence_length, max_training_elements_per_user=max_elements_per_user, max_batches_per_user=max_batches_per_client) input_spec = train_dataset_preprocess_comp.type_signature.result.element def tff_model_fn() -> tff.learning.Model: return tff.learning.from_keras_model( keras_model=model_builder(), input_spec=input_spec, loss=loss_builder(), metrics=metrics_builder()) def client_weight_fn(local_outputs): # Num_tokens is a tensor with type int64[1], to use as a weight need # a float32 scalar. return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32) training_process = iterative_process_builder( tff_model_fn, client_weight_fn=client_weight_fn) training_process = tff.simulation.compose_dataset_computation_with_iterative_process( train_dataset_preprocess_comp, training_process) client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=train_clientdata, train_clients_per_round=clients_per_round, random_seed=client_datasets_random_seed) evaluate_fn = training_utils.build_evaluate_fn( model_builder=model_builder, eval_dataset=validation_dataset, loss_builder=loss_builder, metrics_builder=metrics_builder) test_fn = training_utils.build_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) logging.info('Training model:') logging.info(model_builder().summary()) training_loop.run( iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs)
def run_federated( iterative_process_builder: Callable[..., tff.templates.IterativeProcess], client_epochs_per_round: int, client_batch_size: int, clients_per_round: int, max_batches_per_client: Optional[int] = -1, client_datasets_random_seed: Optional[int] = None, vocab_tokens_size: Optional[int] = 10000, vocab_tags_size: Optional[int] = 500, max_elements_per_user: Optional[int] = 1000, num_validation_examples: Optional[int] = 10000, total_rounds: Optional[int] = 1500, experiment_name: Optional[str] = 'federated_so_lr', root_output_dir: Optional[str] = '/tmp/fed_opt', max_eval_batches: Optional[int] = None, **kwargs): """Runs an iterative process on the Stack Overflow logistic regression task. This method will load and pre-process dataset and construct a model used for the task. It then uses `iterative_process_builder` to create an iterative process that it applies to the task, using `federated_research.utils.training_loop`. We assume that the iterative process has the following functional type signatures: * `initialize`: `( -> S@SERVER)` where `S` represents the server state. * `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S` represents the server state, `{B*}` represents the client datasets, and `T` represents a python `Mapping` object. Moreover, the server state must have an attribute `model` of type `tff.learning.ModelWeights`. Args: iterative_process_builder: A function that accepts a no-arg `model_fn`, and returns a `tff.templates.IterativeProcess`. The `model_fn` must return a `tff.learning.Model`. client_epochs_per_round: An integer representing the number of epochs of training performed per client in each training round. client_batch_size: An integer representing the batch size used on clients. clients_per_round: An integer representing the number of clients participating in each round. max_batches_per_client: An optional int specifying the number of batches taken by each client at each round. If `-1`, the entire client dataset is used. client_datasets_random_seed: An optional int used to seed which clients are sampled at each round. If `None`, no seed is used. vocab_tokens_size: Integer dictating the number of most frequent words to use in the vocabulary. vocab_tags_size: Integer dictating the number of most frequent tags to use in the label creation. max_elements_per_user: The maximum number of elements processed for each client's dataset. num_validation_examples: The number of test examples to use for validation. total_rounds: The number of federated training rounds. experiment_name: The name of the experiment being run. This will be appended to the `root_output_dir` for purposes of writing outputs. root_output_dir: The name of the root output directory for writing experiment outputs. max_eval_batches: If set to a positive integer, evaluation datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full evaluation datasets are used. **kwargs: Additional arguments configuring the training loop. For details on supported arguments, see `federated_research/utils/training_utils.py`. """ stackoverflow_train, _, _ = stackoverflow_lr_dataset.get_stackoverflow_datasets( vocab_tokens_size=vocab_tokens_size, vocab_tags_size=vocab_tags_size, client_batch_size=client_batch_size, client_epochs_per_round=client_epochs_per_round, max_training_elements_per_user=max_elements_per_user, max_batches_per_user=max_batches_per_client, num_validation_examples=num_validation_examples) _, stackoverflow_validation, stackoverflow_test = stackoverflow_lr_dataset.get_centralized_datasets( train_batch_size=client_batch_size, vocab_tokens_size=vocab_tokens_size, vocab_tags_size=vocab_tags_size, num_validation_examples=num_validation_examples, max_validation_batches=max_eval_batches, max_test_batches=max_eval_batches) input_spec = stackoverflow_train.create_tf_dataset_for_client( stackoverflow_train.client_ids[0]).element_spec model_builder = functools.partial( stackoverflow_lr_models.create_logistic_model, vocab_tokens_size=vocab_tokens_size, vocab_tags_size=vocab_tags_size) loss_builder = functools.partial(tf.keras.losses.BinaryCrossentropy, from_logits=False, reduction=tf.keras.losses.Reduction.SUM) def tff_model_fn() -> tff.learning.Model: return tff.learning.from_keras_model(keras_model=model_builder(), input_spec=input_spec, loss=loss_builder(), metrics=metrics_builder()) training_process = iterative_process_builder(tff_model_fn) client_datasets_fn = training_utils.build_client_datasets_fn( train_dataset=stackoverflow_train, train_clients_per_round=clients_per_round, random_seed=client_datasets_random_seed) evaluate_fn = training_utils.build_evaluate_fn( model_builder=model_builder, eval_dataset=stackoverflow_validation, loss_builder=loss_builder, metrics_builder=metrics_builder) test_fn = training_utils.build_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=stackoverflow_validation.concatenate(stackoverflow_test), loss_builder=loss_builder, metrics_builder=metrics_builder) logging.info('Training model:') logging.info(model_builder().summary()) training_loop.run(iterative_process=training_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, test_fn=test_fn, total_rounds=total_rounds, experiment_name=experiment_name, root_output_dir=root_output_dir, **kwargs)
def run_experiment(): """Data preprocessing and experiment execution.""" emnist_train, emnist_test = emnist_dataset.get_emnist_datasets( FLAGS.client_batch_size, FLAGS.client_epochs_per_round, 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 client_datasets_fn = training_utils.build_client_datasets_fn( emnist_train, FLAGS.clients_per_round) evaluate_fn = training_utils.build_evaluate_fn( eval_dataset=emnist_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) 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) hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags()) training_loop_dict = utils_impl.lookup_flag_values(training_loop_flags) training_loop.run( iterative_process=iterative_process, client_datasets_fn=client_datasets_fn, validation_fn=evaluate_fn, hparam_dict=hparam_dict, **training_loop_dict)