def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
# TODO(b/139129100): Remove this once the local executor is the default.
tff.framework.set_default_executor(
tff.framework.local_executor_factory(max_fanout=25))
emnist_train, _ = dataset.get_emnist_datasets(only_digits=False)
emnist_train = emnist_train.preprocess(lambda x: x.batch(20))
input_spec = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0]).element_spec
if FLAGS.model == 'cnn':
model_builder = functools.partial(models.create_conv_dropout_model,
only_digits=False)
elif FLAGS.model == '2nn':
model_builder = functools.partial(models.create_two_hidden_layer_model,
only_digits=False)
else:
raise ValueError('Cannot handle model flag [{!s}].'.format(
FLAGS.model))
tff_model = tff.learning.from_keras_model(
keras_model=model_builder(),
input_spec=input_spec,
loss=tf.keras.losses.SparseCategoricalCrossentropy())
yogi_init_accum_estimate = optimizer_utils.compute_yogi_init(
emnist_train, tff_model, num_clients=FLAGS.num_clients)
logging.info('Yogi initializer: {:s}'.format(
format(yogi_init_accum_estimate, '10.6E')))
def run_experiment():
"""Data preprocessing and experiment execution."""
emnist_train, emnist_test = 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])
sample_batch = tf.nest.map_structure(lambda x: x.numpy(),
next(iter(example_dataset)))
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,
dummy_batch=sample_batch,
loss=loss_builder(),
metrics=metrics_builder())
if FLAGS.use_compression:
# We create a `StatefulBroadcastFn` and `StatefulAggregateFn` 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_fn = (
tff.learning.framework.build_encoded_broadcast_from_model(
tff_model_fn, _broadcast_encoder_fn))
encoded_mean_fn = tff.learning.framework.build_encoded_mean_from_model(
tff_model_fn, _mean_encoder_fn)
else:
encoded_broadcast_fn = None
encoded_mean_fn = 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,
stateful_delta_aggregate_fn=encoded_mean_fn,
stateful_model_broadcast_fn=encoded_broadcast_fn)
iterative_process = compression_process_adapter.CompressionProcessAdapter(
iterative_process)
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
evaluate_fn=evaluate_fn)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
tf.compat.v1.enable_v2_behavior()
# TODO(b/139129100): Remove this once the local executor is the default.
tff.framework.set_default_executor(
tff.framework.local_executor_factory(max_fanout=25))
emnist_train, emnist_test = dataset.get_emnist_datasets(
FLAGS.client_batch_size,
FLAGS.client_epochs_per_round,
only_digits=False)
sample_client_dataset = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0])
# TODO(b/144382142): Sample batches cannot be eager tensors, since they are
# passed (implicitly) to tff.learning.build_federated_averaging_process.
sample_batch = tf.nest.map_structure(lambda x: x.numpy(),
next(iter(sample_client_dataset)))
if FLAGS.model == 'cnn':
model_builder = functools.partial(models.create_conv_dropout_model,
only_digits=False)
elif FLAGS.model == '2nn':
model_builder = functools.partial(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()]
training_process = iterative_process_builder.from_flags(
dummy_batch=sample_batch,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
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())
training_loop.run(
iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
evaluate_fn=evaluate_fn,
)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
tf.compat.v1.enable_v2_behavior()
emnist_train, emnist_test = dataset.get_emnist_datasets(
FLAGS.client_batch_size,
FLAGS.client_epochs_per_round,
only_digits=False)
input_spec = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0]).element_spec
if FLAGS.model == 'cnn':
model_builder = functools.partial(models.create_conv_dropout_model,
only_digits=False)
elif FLAGS.model == '2nn':
model_builder = functools.partial(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()]
training_process = iterative_process_builder.from_flags(
input_spec=input_spec,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset=emnist_train,
train_clients_per_round=FLAGS.clients_per_round,
random_seed=FLAGS.client_datasets_random_seed)
assign_weights_fn = fed_avg_schedule.ServerState.assign_weights_to_keras_model
evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=emnist_test,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder,
assign_weights_to_keras_model=assign_weights_fn)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(
iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
evaluate_fn=evaluate_fn,
)
def test_take_with_repeat(self):
emnist_train, _ = dataset.get_emnist_datasets(
client_batch_size=10,
client_epochs_per_round=-1,
max_batches_per_client=10,
only_digits=True)
self.assertEqual(len(emnist_train.client_ids), 3383)
for i in range(10):
client_ds = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[i])
self.assertEqual(_compute_length_of_dataset(client_ds), 10)
def test_emnist_dataset_structure(self):
emnist_train, emnist_test = dataset.get_emnist_datasets(
client_batch_size=10, client_epochs_per_round=1, only_digits=True)
self.assertEqual(len(emnist_train.client_ids), 3383)
sample_train_ds = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0])
train_batch = next(iter(sample_train_ds))
train_batch_shape = train_batch[0].shape
test_batch = next(iter(emnist_test))
test_batch_shape = test_batch[0].shape
self.assertEqual(train_batch_shape.as_list(), [10, 28, 28, 1])
self.assertEqual(test_batch_shape.as_list(), [TEST_BATCH_SIZE, 28, 28, 1])
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
tf.compat.v1.enable_v2_behavior()
# TODO(b/139129100): Remove this once the local executor is the default.
tff.framework.set_default_executor(
tff.framework.local_executor_factory(max_fanout=25))
emnist_train, _ = dataset.get_emnist_datasets(only_digits=False)
emnist_train = emnist_train.preprocess(lambda x: x.batch(20))
sample_client_dataset = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0])
# TODO(b/144382142): Sample batches cannot be eager tensors, since they are
# passed (implicitly) to tff.learning.build_federated_averaging_process.
sample_batch = tf.nest.map_structure(lambda x: x.numpy(),
next(iter(sample_client_dataset)))
if FLAGS.model == 'cnn':
model_builder = functools.partial(models.create_conv_dropout_model,
only_digits=False)
elif FLAGS.model == '2nn':
model_builder = functools.partial(models.create_two_hidden_layer_model,
only_digits=False)
else:
raise ValueError('Cannot handle model flag [{!s}].'.format(
FLAGS.model))
tff_model = tff.learning.from_keras_model(
keras_model=model_builder(),
dummy_batch=sample_batch,
loss=tf.keras.losses.SparseCategoricalCrossentropy())
yogi_init_accum_estimate = optimizer_utils.compute_yogi_init(
emnist_train, tff_model, num_clients=FLAGS.num_clients)
logging.info('Yogi initializer: {:s}'.format(
format(yogi_init_accum_estimate, '10.6E')))
def test_raises_no_repeat_and_no_take(self):
with self.assertRaises(ValueError):
dataset.get_emnist_datasets(
client_batch_size=10,
client_epochs_per_round=-1,
max_batches_per_client=-1)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
tf.compat.v1.enable_v2_behavior()
emnist_train, emnist_test = dataset.get_emnist_datasets(
FLAGS.client_batch_size, FLAGS.client_epochs_per_round, only_digits=False)
if FLAGS.model == 'cnn':
model_builder = functools.partial(
models.create_conv_dropout_model, only_digits=False)
elif FLAGS.model == '2nn':
model_builder = functools.partial(
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_num_clients=FLAGS.clients_per_round,
per_vector_clipping=FLAGS.per_vector_clipping,
model=model_fn())
dp_aggregate_fn, _ = tff.utils.build_dp_aggregate(dp_query)
else:
dp_aggregate_fn = None
server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('server')
client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('client')
training_process = (
tff.learning.federated_averaging.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,
stateful_delta_aggregate_fn=dp_aggregate_fn))
adaptive_clipping = (FLAGS.adaptive_clip_learning_rate > 0)
training_process = dp_utils.DPFedAvgProcessAdapter(training_process,
FLAGS.per_vector_clipping,
adaptive_clipping)
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,
assign_weights_to_keras_model=dp_utils.assign_weights_to_keras_model)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(
iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
evaluate_fn=evaluate_fn,
)
