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