def test_preprocess_applied(self, mock_load_data):
if tf.config.list_logical_devices('GPU'):
self.skipTest('skip GPU test')
# Mock out the actual data loading from disk. Assert that the preprocessing
# function is applied to the client data, and that only the ClientData
# objects we desired are used.
#
# The correctness of the preprocessing function is tested in other tests.
sample_ds = tf.data.Dataset.from_tensor_slices(TEST_DATA)
mock_train = mock.create_autospec(tff.simulation.datasets.ClientData)
mock_train.create_tf_dataset_from_all_clients = mock.Mock(
return_value=sample_ds)
mock_test = mock.create_autospec(tff.simulation.datasets.ClientData)
mock_test.create_tf_dataset_from_all_clients = mock.Mock(
return_value=sample_ds)
mock_load_data.return_value = (mock_train, mock_test)
_, _ = emnist_dataset.get_centralized_datasets()
mock_load_data.assert_called_once()
# Assert the validation ClientData isn't used, and the train and test
# are amalgamated into datasets single datasets over all clients.
self.assertEqual(mock_train.mock_calls,
mock.call.create_tf_dataset_from_all_clients().call_list())
self.assertEqual(mock_test.mock_calls,
mock.call.create_tf_dataset_from_all_clients().call_list())
def test_global_emnist_dataset_structure(self):
global_train, global_test = emnist_dataset.get_centralized_datasets(
train_batch_size=32, test_batch_size=100, only_digits=False)
train_batch = next(iter(global_train))
train_batch_shape = train_batch[0].shape
test_batch = next(iter(global_test))
test_batch_shape = test_batch[0].shape
self.assertEqual(train_batch_shape.as_list(), [32, 28, 28, 1])
self.assertEqual(test_batch_shape.as_list(), [100, 28, 28, 1])
def run_centralized(optimizer: tf.keras.optimizers.Optimizer,
experiment_name: str,
root_output_dir: str,
num_epochs: int,
batch_size: int,
decay_epochs: Optional[int] = None,
lr_decay: Optional[float] = None,
hparams_dict: Optional[Mapping[str, Any]] = None,
max_batches: Optional[int] = None,
cache_dir: Optional[str] = None):
"""Trains a bottleneck autoencoder on EMNIST using a given optimizer.
Args:
optimizer: A `tf.keras.optimizers.Optimizer` used to perform training.
experiment_name: The name of the experiment. Part of the output directory.
root_output_dir: The top-level output directory for experiment runs. The
`experiment_name` argument will be appended, and the directory will
contain tensorboard logs, metrics written as CSVs, and a CSV of
hyperparameter choices (if `hparams_dict` is used).
num_epochs: The number of training epochs.
batch_size: The batch size, used for train, validation, and test.
decay_epochs: The number of epochs of training before decaying the learning
rate. If None, no decay occurs.
lr_decay: The amount to decay the learning rate by after `decay_epochs`
training epochs have occurred.
hparams_dict: A mapping with string keys representing the hyperparameters
and their values. If not None, this is written to CSV.
max_batches: If set to a positive integer, datasets are capped to at most
that many batches. If set to None or a nonpositive integer, the full
datasets are used.
"""
train_dataset, eval_dataset = emnist_dataset.get_centralized_datasets(
train_batch_size=batch_size, only_digits=False, emnist_task='autoencoder', cache_dir=cache_dir)
if max_batches and max_batches >= 1:
train_dataset = train_dataset.take(max_batches)
eval_dataset = eval_dataset.take(max_batches)
model = emnist_ae_models.create_autoencoder_model()
model.compile(
loss=tf.keras.losses.MeanSquaredError(),
optimizer=optimizer,
metrics=[tf.keras.metrics.MeanSquaredError()])
centralized_training_loop.run(
keras_model=model,
train_dataset=train_dataset,
validation_dataset=eval_dataset,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
num_epochs=num_epochs,
hparams_dict=hparams_dict,
decay_epochs=decay_epochs,
lr_decay=lr_decay)
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 configure_training(
task_spec: training_specs.TaskSpec) -> training_specs.RunnerSpec:
"""Configures training for the EMNIST autoencoder task.
This method will load and pre-process datasets and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process compatible with `federated_research.utils.training_loop`.
Args:
task_spec: A `TaskSpec` class for creating federated training tasks.
Returns:
A `RunnerSpec` containing attributes used for running the newly created
federated task.
"""
emnist_task = 'autoencoder'
emnist_train, _ = tff.simulation.datasets.emnist.load_data(only_digits=False)
_, emnist_test = emnist_dataset.get_centralized_datasets(
only_digits=False, emnist_task=emnist_task)
train_preprocess_fn = emnist_dataset.create_preprocess_fn(
num_epochs=task_spec.client_epochs_per_round,
batch_size=task_spec.client_batch_size,
emnist_task=emnist_task)
input_spec = train_preprocess_fn.type_signature.result.element
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())
iterative_process = task_spec.iterative_process_builder(tff_model_fn)
if hasattr(emnist_train, 'dataset_computation'):
@tff.tf_computation(tf.string)
def build_train_dataset_from_client_id(client_id):
client_dataset = emnist_train.dataset_computation(client_id)
return train_preprocess_fn(client_dataset)
training_process = tff.simulation.compose_dataset_computation_with_iterative_process(
build_train_dataset_from_client_id, iterative_process)
client_ids_fn = training_utils.build_sample_fn(
emnist_train.client_ids,
size=task_spec.clients_per_round,
replace=False,
random_seed=task_spec.sampling_random_seed)
# We convert the output to a list (instead of an np.ndarray) so that it can
# be used as input to the iterative process.
client_sampling_fn = lambda x: list(client_ids_fn(x))
else:
training_process = tff.simulation.compose_dataset_computation_with_iterative_process(
train_preprocess_fn, iterative_process)
client_sampling_fn = training_utils.build_client_datasets_fn(
dataset=emnist_train,
clients_per_round=task_spec.clients_per_round,
random_seed=task_spec.sampling_random_seed)
training_process.get_model_weights = iterative_process.get_model_weights
test_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: test_fn(model_weights)
return training_specs.RunnerSpec(
iterative_process=training_process,
client_datasets_fn=client_sampling_fn,
validation_fn=validation_fn,
test_fn=test_fn)
def configure_training(task_spec: training_specs.TaskSpec,
model: str = 'cnn') -> training_specs.RunnerSpec:
"""Configures training for the EMNIST character recognition task.
This method will load and pre-process datasets and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process compatible with `federated_research.utils.training_loop`.
Args:
task_spec: A `TaskSpec` class for creating federated training tasks.
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).
Returns:
A `RunnerSpec` containing attributes used for running the newly created
federated task.
"""
emnist_task = 'digit_recognition'
emnist_train, _ = tff.simulation.datasets.emnist.load_data(only_digits=False)
_, emnist_test = emnist_dataset.get_centralized_datasets(
only_digits=False, emnist_task=emnist_task)
train_preprocess_fn = emnist_dataset.create_preprocess_fn(
num_epochs=task_spec.client_epochs_per_round,
batch_size=task_spec.client_batch_size,
emnist_task=emnist_task)
input_spec = train_preprocess_fn.type_signature.result.element
if model == 'cnn':
model_builder = functools.partial(
emnist_models.create_conv_dropout_model, only_digits=False)
elif model == '2nn':
model_builder = functools.partial(
emnist_models.create_two_hidden_layer_model, only_digits=False)
else:
raise ValueError(
'Cannot handle model flag [{!s}], must be one of {!s}.'.format(
model, EMNIST_MODELS))
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())
iterative_process = task_spec.iterative_process_builder(tff_model_fn)
@tff.tf_computation(tf.string)
def build_train_dataset_from_client_id(client_id):
client_dataset = emnist_train.dataset_computation(client_id)
return train_preprocess_fn(client_dataset)
training_process = tff.simulation.compose_dataset_computation_with_iterative_process(
build_train_dataset_from_client_id, iterative_process)
client_ids_fn = tff.simulation.build_uniform_sampling_fn(
emnist_train.client_ids,
size=task_spec.clients_per_round,
replace=False,
random_seed=task_spec.client_datasets_random_seed)
# We convert the output to a list (instead of an np.ndarray) so that it can
# be used as input to the iterative process.
client_sampling_fn = lambda x: list(client_ids_fn(x))
training_process.get_model_weights = iterative_process.get_model_weights
evaluate_fn = tff.learning.build_federated_evaluation(tff_model_fn)
def test_fn(state):
return evaluate_fn(
iterative_process.get_model_weights(state), [emnist_test])
def validation_fn(state, round_num):
del round_num
return evaluate_fn(
iterative_process.get_model_weights(state), [emnist_test])
return training_specs.RunnerSpec(
iterative_process=training_process,
client_datasets_fn=client_sampling_fn,
validation_fn=validation_fn,
test_fn=test_fn)
def run_federated(
iterative_process_builder: Callable[..., tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
client_datasets_random_seed: Optional[int] = None,
model: Optional[str] = 'cnn',
total_rounds: Optional[int] = 1500,
experiment_name: Optional[str] = 'federated_emnist_cr',
root_output_dir: Optional[str] = '/tmp/fed_opt',
**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.
The iterative process must also have a callable attribute `get_model_weights`
that takes as input the state of the iterative process, and returns a
`tff.learning.ModelWeights` object.
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.
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.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`federated_research/utils/training_utils.py`.
"""
emnist_train, _ = emnist_dataset.get_federated_datasets(
train_client_batch_size=client_batch_size,
train_client_epochs_per_round=client_epochs_per_round,
only_digits=False)
_, emnist_test = emnist_dataset.get_centralized_datasets(only_digits=False)
input_spec = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0]).element_spec
if model == 'cnn':
model_builder = functools.partial(
emnist_models.create_conv_dropout_model, only_digits=False)
elif model == '2nn':
model_builder = functools.partial(
emnist_models.create_two_hidden_layer_model, only_digits=False)
else:
raise ValueError(
'Cannot handle model flag [{!s}], must be one of {!s}.'.format(
model, EMNIST_MODELS))
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)
client_datasets_fn = training_utils.build_client_datasets_fn(
dataset=emnist_train,
clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed)
test_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: test_fn(model_weights)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(
iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
test_fn=test_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))
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 main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
emnist_task = 'digit_recognition'
emnist_train, _ = tff.simulation.datasets.emnist.load_data(only_digits=False)
_, emnist_test = emnist_dataset.get_centralized_datasets(
only_digits=False, emnist_task=emnist_task)
train_preprocess_fn = emnist_dataset.create_preprocess_fn(
num_epochs=FLAGS.client_epochs_per_round,
batch_size=FLAGS.client_batch_size,
emnist_task=emnist_task)
input_spec = train_preprocess_fn.type_signature.result.element
if FLAGS.model == 'cnn':
model_builder = functools.partial(
emnist_models.create_conv_dropout_model, only_digits=FLAGS.only_digits)
elif FLAGS.model == '2nn':
model_builder = functools.partial(
emnist_models.create_two_hidden_layer_model,
only_digits=FLAGS.only_digits)
elif FLAGS.model == '1m_cnn':
model_builder = functools.partial(
create_1m_cnn_model, only_digits=FLAGS.only_digits)
else:
raise ValueError('Cannot handle model flag [{!s}].'.format(FLAGS.model))
logging.info('Training model:')
logging.info(model_builder().summary())
loss_builder = tf.keras.losses.SparseCategoricalCrossentropy
metrics_builder = lambda: [tf.keras.metrics.SparseCategoricalAccuracy()]
compression_dict = utils_impl.lookup_flag_values(compression_flags)
dp_dict = utils_impl.lookup_flag_values(dp_flags)
# Most logic for deciding what baseline to run is here.
aggregation_factory = fl_utils.build_aggregator(
compression_flags=compression_dict,
dp_flags=dp_dict,
num_clients=len(emnist_train.client_ids),
num_clients_per_round=FLAGS.clients_per_round,
num_rounds=FLAGS.total_rounds,
client_template=model_builder().trainable_variables)
def tff_model_fn():
return tff.learning.from_keras_model(
keras_model=model_builder(),
loss=loss_builder(),
input_spec=input_spec,
metrics=metrics_builder())
server_optimizer_fn = lambda: utils_impl.create_optimizer_from_flags('server')
client_optimizer_fn = lambda: utils_impl.create_optimizer_from_flags('client')
iterative_process = tff.learning.build_federated_averaging_process(
model_fn=tff_model_fn,
server_optimizer_fn=server_optimizer_fn,
client_weighting=tff.learning.ClientWeighting.UNIFORM,
client_optimizer_fn=client_optimizer_fn,
model_update_aggregation_factory=aggregation_factory)
@tff.tf_computation(tf.string)
def build_train_dataset_from_client_id(client_id):
client_dataset = emnist_train.dataset_computation(client_id)
return train_preprocess_fn(client_dataset)
training_process = tff.simulation.compose_dataset_computation_with_iterative_process(
build_train_dataset_from_client_id, iterative_process)
training_process.get_model_weights = iterative_process.get_model_weights
client_ids_fn = functools.partial(
tff.simulation.build_uniform_sampling_fn(
emnist_train.client_ids,
replace=False,
random_seed=FLAGS.client_datasets_random_seed),
size=FLAGS.clients_per_round)
# We convert the output to a list (instead of an np.ndarray) so that it can
# be used as input to the iterative process.
client_sampling_fn = lambda x: list(client_ids_fn(x))
evaluate_fn = tff.learning.build_federated_evaluation(tff_model_fn)
def test_fn(state):
return evaluate_fn(
iterative_process.get_model_weights(state), [emnist_test])
def validation_fn(state, round_num):
del round_num
return evaluate_fn(
iterative_process.get_model_weights(state), [emnist_test])
training_loop.run(
iterative_process=training_process,
client_datasets_fn=client_sampling_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 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:
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)
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_weighting=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_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())
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=validation_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_emnist_cr',
root_output_dir: Optional[str] = '/tmp/fed_opt',
max_eval_batches: Optional[int] = None,
**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`.
"""
emnist_train, _, federated_test = emnist_dataset.get_emnist_datasets(
client_batch_size,
client_epochs_per_round,
max_batches_per_client=max_batches_per_client,
only_digits=False)
_, emnist_test = emnist_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
if model == 'cnn':
model_builder = functools.partial(
emnist_models.create_conv_dropout_model, only_digits=False)
elif model == '2nn':
model_builder = functools.partial(
emnist_models.create_two_hidden_layer_model, only_digits=False)
else:
raise ValueError(
'Cannot handle model flag [{!s}], must be one of {!s}.'.format(
model, EMNIST_MODELS))
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(model_fn=tff_model_fn)
evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=emnist_test,
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']
q_client = np.load(
f'/home/monica/AVAIL_VECTORS/q_client_{var}_emnist.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=emnist_train,
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=emnist_train,
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']
client_datasets_fn = training_utils.build_availability_client_datasets_fn(
train_dataset=emnist_train,
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_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)
def run_centralized(optimizer: tf.keras.optimizers.Optimizer,
experiment_name: str,
root_output_dir: str,
num_epochs: int,
batch_size: int,
decay_epochs: Optional[int] = None,
lr_decay: Optional[float] = None,
hparams_dict: Optional[Mapping[str, Any]] = None,
emnist_model: Optional[str] = 'cnn',
max_batches: Optional[int] = None):
"""Trains a model on EMNIST character recognition using a given optimizer.
Args:
optimizer: A `tf.keras.optimizers.Optimizer` used to perform training.
experiment_name: The name of the experiment. Part of the output directory.
root_output_dir: The top-level output directory for experiment runs. The
`experiment_name` argument will be appended, and the directory will
contain tensorboard logs, metrics written as CSVs, and a CSV of
hyperparameter choices (if `hparams_dict` is used).
num_epochs: The number of training epochs.
batch_size: The batch size, used for train, validation, and test.
decay_epochs: The number of epochs of training before decaying the learning
rate. If None, no decay occurs.
lr_decay: The amount to decay the learning rate by after `decay_epochs`
training epochs have occurred.
hparams_dict: A mapping with string keys representing the hyperparameters
and their values. If not None, this is written to CSV.
emnist_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).
max_batches: If set to a positive integer, datasets are capped to at most
that many batches. If set to None or a nonpositive integer, the full
datasets are used.
"""
train_dataset, eval_dataset = emnist_dataset.get_centralized_datasets(
train_batch_size=batch_size,
max_train_batches=max_batches,
max_test_batches=max_batches,
only_digits=False)
if emnist_model == 'cnn':
model = emnist_models.create_conv_dropout_model(only_digits=False)
elif emnist_model == '2nn':
model = emnist_models.create_two_hidden_layer_model(only_digits=False)
else:
raise ValueError(
'Cannot handle model flag [{!s}].'.format(emnist_model))
model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
optimizer=optimizer,
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
centralized_training_loop.run(keras_model=model,
train_dataset=train_dataset,
validation_dataset=eval_dataset,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
num_epochs=num_epochs,
hparams_dict=hparams_dict,
decay_epochs=decay_epochs,
lr_decay=lr_decay)
def configure_training(task_spec: training_specs.TaskSpec,
eval_spec: Optional[training_specs.EvalSpec] = None,
model: str = 'cnn') -> training_specs.RunnerSpec:
"""Configures training for the EMNIST character recognition task.
This method will load and pre-process datasets and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process compatible with `federated_research.utils.training_loop`.
Args:
task_spec: A `TaskSpec` class for creating federated training tasks.
eval_spec: An `EvalSpec` class for configuring federated evaluation. If set
to None, centralized evaluation is used for validation and testing
instead.
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).
Returns:
A `RunnerSpec` containing attributes used for running the newly created
federated task.
"""
emnist_task = 'digit_recognition'
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(
only_digits=False)
train_preprocess_fn = emnist_dataset.create_preprocess_fn(
num_epochs=task_spec.client_epochs_per_round,
batch_size=task_spec.client_batch_size,
emnist_task=emnist_task)
input_spec = train_preprocess_fn.type_signature.result.element
if model == 'cnn':
model_builder = functools.partial(
emnist_models.create_conv_dropout_model, only_digits=False)
elif model == '2nn':
model_builder = functools.partial(
emnist_models.create_two_hidden_layer_model, only_digits=False)
else:
raise ValueError(
'Cannot handle model flag [{!s}], must be one of {!s}.'.format(
model, EMNIST_MODELS))
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())
iterative_process = task_spec.iterative_process_builder(tff_model_fn)
clients_per_train_round = min(task_spec.clients_per_round,
TOTAL_NUM_TRAIN_CLIENTS)
if hasattr(emnist_train, 'dataset_computation'):
@tff.tf_computation(tf.string)
def build_train_dataset_from_client_id(client_id):
client_dataset = emnist_train.dataset_computation(client_id)
return train_preprocess_fn(client_dataset)
training_process = tff.simulation.compose_dataset_computation_with_iterative_process(
build_train_dataset_from_client_id, iterative_process)
client_ids_fn = training_utils.build_sample_fn(
emnist_train.client_ids,
size=clients_per_train_round,
replace=False,
random_seed=task_spec.sampling_random_seed)
# We convert the output to a list (instead of an np.ndarray) so that it can
# be used as input to the iterative process.
client_sampling_fn = lambda x: list(client_ids_fn(x))
else:
training_process = tff.simulation.compose_dataset_computation_with_iterative_process(
train_preprocess_fn, iterative_process)
client_sampling_fn = training_utils.build_client_datasets_fn(
dataset=emnist_train,
clients_per_round=clients_per_train_round,
random_seed=task_spec.sampling_random_seed)
training_process.get_model_weights = iterative_process.get_model_weights
if eval_spec:
if eval_spec.clients_per_validation_round is None:
clients_per_validation_round = TOTAL_NUM_TEST_CLIENTS
else:
clients_per_validation_round = min(eval_spec.clients_per_validation_round,
TOTAL_NUM_TEST_CLIENTS)
if eval_spec.clients_per_test_round is None:
clients_per_test_round = TOTAL_NUM_TEST_CLIENTS
else:
clients_per_test_round = min(eval_spec.clients_per_test_round,
TOTAL_NUM_TEST_CLIENTS)
test_preprocess_fn = emnist_dataset.create_preprocess_fn(
num_epochs=1,
batch_size=eval_spec.client_batch_size,
shuffle_buffer_size=1,
emnist_task=emnist_task)
emnist_test = emnist_test.preprocess(test_preprocess_fn)
def eval_metrics_builder():
return [
tf.keras.metrics.SparseCategoricalCrossentropy(),
tf.keras.metrics.SparseCategoricalAccuracy()
]
federated_eval_fn = training_utils.build_federated_evaluate_fn(
model_builder=model_builder, metrics_builder=eval_metrics_builder)
validation_client_sampling_fn = training_utils.build_client_datasets_fn(
emnist_test,
clients_per_validation_round,
random_seed=eval_spec.sampling_random_seed)
test_client_sampling_fn = training_utils.build_client_datasets_fn(
emnist_test,
clients_per_test_round,
random_seed=eval_spec.sampling_random_seed)
def validation_fn(model_weights, round_num):
validation_clients = validation_client_sampling_fn(round_num)
return federated_eval_fn(model_weights, validation_clients)
def test_fn(model_weights):
# We fix the round number to get deterministic behavior
test_round_num = 0
test_clients = test_client_sampling_fn(test_round_num)
return federated_eval_fn(model_weights, test_clients)
else:
_, central_emnist_test = emnist_dataset.get_centralized_datasets(
only_digits=False, emnist_task=emnist_task)
test_fn = training_utils.build_centralized_evaluate_fn(
eval_dataset=central_emnist_test,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
validation_fn = lambda model_weights, round_num: test_fn(model_weights)
return training_specs.RunnerSpec(
iterative_process=training_process,
client_datasets_fn=client_sampling_fn,
validation_fn=validation_fn,
test_fn=test_fn)
def run_federated(
iterative_process_builder: Callable[...,
tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
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.
The iterative process must also have a callable attribute `get_model_weights`
that takes as input the state of the iterative process, and returns a
`tff.learning.ModelWeights` object.
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.
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_dataset.get_federated_datasets(
train_client_batch_size=client_batch_size,
train_client_epochs_per_round=client_epochs_per_round,
only_digits=False,
emnist_task='autoencoder')
_, emnist_test = emnist_dataset.get_centralized_datasets(
train_batch_size=client_batch_size,
only_digits=False,
emnist_task='autoencoder')
if max_eval_batches and max_eval_batches >= 1:
emnist_test = emnist_test.take(max_eval_batches)
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(
dataset=emnist_train,
clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed)
evaluate_fn = training_utils.build_centralized_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)
