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,
sequence_length: Optional[int] = 80):
"""Trains a two-layer RNN on Shakespeare next-character-prediction.
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.
sequence_length: The sequence length used for Shakespeare preprocessing.
"""
train_dataset, eval_dataset = shakespeare_dataset.construct_centralized_datasets(
batch_size=batch_size, sequence_length=sequence_length)
pad_token, _, _, _ = shakespeare_dataset.get_special_tokens()
model = shakespeare_models.create_recurrent_model(
vocab_size=VOCAB_SIZE, sequence_length=sequence_length)
model.compile(
optimizer=optimizer,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[
keras_metrics.MaskedCategoricalAccuracy(masked_tokens=[pad_token])
])
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 main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
train_clientdata = shakespeare_dataset.construct_character_level_datasets(
FLAGS.client_batch_size, FLAGS.client_epochs_per_round,
FLAGS.sequence_length)
_, test_dataset = shakespeare_dataset.construct_centralized_datasets()
test_dataset = test_dataset.cache()
loss_fn_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)
training_process = iterative_process_builder.from_flags(
input_spec=input_spec,
model_builder=model_builder,
loss_builder=loss_fn_builder,
metrics_builder=metrics_builder,
client_weight_fn=client_weight_fn)
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset=train_clientdata,
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=test_dataset,
model_builder=model_builder,
loss_builder=loss_fn_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,
validation_fn=evaluate_fn)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
train_clientdata = shakespeare_dataset.construct_character_level_datasets(
FLAGS.client_batch_size,
FLAGS.client_epochs_per_round,
sequence_length=FLAGS.sequence_length,
max_batches_per_client=FLAGS.max_batches_per_client,
shuffle_buffer_size=0)
eval_train_dataset, eval_test_dataset = (
shakespeare_dataset.construct_centralized_datasets())
loss_fn_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
if FLAGS.client_weight == 'uniform':
def client_weight_fn(local_outputs):
del local_outputs
return 1.0
elif FLAGS.client_weight == 'num_tokens':
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)
else:
raise ValueError(
'Unsupported client_weight flag [{!s}]. Currently only '
'`uniform` and `num_tokens` are supported.'.format(
FLAGS.client_weight))
training_process = decay_iterative_process_builder.from_flags(
input_spec=input_spec,
model_builder=model_builder,
loss_builder=loss_fn_builder,
metrics_builder=metrics_builder,
client_weight_fn=client_weight_fn)
client_datasets_fn = training_utils.build_client_datasets_fn(
train_clientdata,
FLAGS.clients_per_round,
random_seed=FLAGS.client_datasets_random_seed)
assign_weights_fn = adaptive_fed_avg.ServerState.assign_weights_to_keras_model
evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=eval_test_dataset,
model_builder=model_builder,
loss_builder=loss_fn_builder,
metrics_builder=metrics_builder,
assign_weights_to_keras_model=assign_weights_fn)
train_evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=eval_train_dataset,
model_builder=model_builder,
loss_builder=loss_fn_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,
validation_fn=evaluate_fn,
train_eval_fn=train_evaluate_fn,
)
def run_federated(
iterative_process_builder: Callable[...,
tff.templates.IterativeProcess],
assign_weights_fn: Callable[[Any, tf.keras.Model], None],
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',
**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
`tensorflow_federated.python.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.
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`.
assign_weights_fn: A function that accepts the server state `S` and a
`tf.keras.Model`, and updates the weights in the Keras model. This is used
to do evaluation using Keras.
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.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`tensorflow_federated/python/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.construct_centralized_datasets()
test_dataset = test_dataset.cache()
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,
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,
validation_fn=evaluate_fn,
test_fn=evaluate_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)