def test_raises_non_callable_get_model_weights_attribute(self):
class BadIterativeProcess(tff.templates.IterativeProcess):
def __init__(self):
pass
def initialize(self):
return {}
def next(self, state, data):
return {}
iterative_process = BadIterativeProcess()
iterative_process.get_model_weights = 2
federated_data = [[_batch_fn()]]
def client_datasets_fn(round_num):
del round_num
return federated_data
def validation_fn(model, round_num):
del model, round_num
return {}
with self.assertRaisesRegex(
training_loop.IterativeProcessCompatibilityError,
_COMPATIBILITY_ERROR_MESSAGE):
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=10,
experiment_name='bad_iterative_process')
def test_fn_writes_metrics(self):
experiment_name = 'test_metrics'
iterative_process = _build_federated_averaging_process()
batch = _batch_fn()
federated_data = [[batch]]
def client_datasets_fn(round_num):
del round_num
return federated_data
def evaluate(model):
keras_model = tff.simulation.models.mnist.create_keras_model(
compile_model=True)
model.assign_weights_to(keras_model)
return {'loss': keras_model.evaluate(batch.x, batch.y)}
root_output_dir = self.get_temp_dir()
training_loop.run(
iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate,
total_rounds=1,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
rounds_per_eval=10,
test_fn=evaluate)
results_dir = os.path.join(root_output_dir, 'results', experiment_name)
scalar_manager = metrics_manager.ScalarMetricsManager(results_dir)
fieldnames, metrics = scalar_manager.get_metrics()
self.assertLen(metrics, 2)
self.assertIn('eval/loss', fieldnames)
self.assertIn('test/loss', fieldnames)
def test_raises_no_model_attribute_in_state(self):
class BadIterativeProcess(tff.templates.IterativeProcess):
"""Converts iterative process results from anonymous tuples."""
def __init__(self):
pass
def initialize(self):
return {}
def next(self, state, data):
return {}
iterative_process = BadIterativeProcess()
federated_data = [[_batch_fn()]]
def client_datasets_fn(round_num):
del round_num
return federated_data
def validation_fn(model):
del model
return {}
with self.assertRaisesRegex(TypeError,
'The server state must have a model attribute'):
training_loop.run(
iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=10,
experiment_name='bad_iterative_process')
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
client_optimizer_fn = lambda: utils_impl.create_optimizer_from_flags(
'client')
server_optimizer_fn = lambda: utils_impl.create_optimizer_from_flags(
'server')
compression_dict = utils_impl.lookup_flag_values(compression_flags)
dp_dict = utils_impl.lookup_flag_values(dp_flags)
def iterative_process_builder(
model_fn: Callable[[], tff.learning.Model],
) -> tff.templates.IterativeProcess:
"""Creates an iterative process using a given TFF `model_fn`."""
model_trainable_variables = model_fn().trainable_variables
# Most logic for deciding what to run is here.
aggregation_factory = fl_utils.build_aggregator(
compression_flags=compression_dict,
dp_flags=dp_dict,
num_clients=get_total_num_clients(FLAGS.task),
num_clients_per_round=FLAGS.clients_per_round,
num_rounds=FLAGS.total_rounds,
client_template=model_trainable_variables)
return tff.learning.build_federated_averaging_process(
model_fn=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)
task_spec = training_specs.TaskSpec(
iterative_process_builder=iterative_process_builder,
client_epochs_per_round=FLAGS.client_epochs_per_round,
client_batch_size=FLAGS.client_batch_size,
clients_per_round=FLAGS.clients_per_round,
client_datasets_random_seed=FLAGS.client_datasets_random_seed)
if FLAGS.task == 'stackoverflow_lr':
runner_spec = federated_stackoverflow_lr.configure_training(task_spec)
else:
raise ValueError(
'--task flag {} is not supported, must be one of {}.'.format(
FLAGS.task, _SUPPORTED_TASKS))
training_loop.run(iterative_process=runner_spec.iterative_process,
client_datasets_fn=runner_spec.client_datasets_fn,
validation_fn=runner_spec.validation_fn,
test_fn=runner_spec.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 test_run_federated(self, task_name, config_fn):
task_spec = training_specs.TaskSpec(
iterative_process_builder=iterative_process_builder,
client_epochs_per_round=1,
client_batch_size=10,
clients_per_round=1,
sampling_random_seed=1)
runner_spec = config_fn(task_spec)
total_rounds = 1
root_output_dir = self.get_temp_dir()
exp_name = 'test_run_federated'
training_loop.run(
iterative_process=runner_spec.iterative_process,
client_datasets_fn=runner_spec.client_datasets_fn,
validation_fn=runner_spec.validation_fn,
# For efficiency, we avoid using the entire test set here
test_fn=None,
total_rounds=total_rounds,
root_output_dir=root_output_dir,
experiment_name=exp_name)
results_dir = os.path.join(root_output_dir, 'results', exp_name)
self.assertTrue(tf.io.gfile.exists(results_dir))
scalar_manager = tff.simulation.CSVMetricsManager(
os.path.join(results_dir, 'experiment.metrics.csv'))
fieldnames, metrics = scalar_manager.get_metrics()
self.assertIn(
'train/train/loss',
fieldnames,
msg=
'The output metrics should have a `train/loss` column if training '
'is successful.')
self.assertIn(
'eval/loss',
fieldnames,
msg=
'The output metrics should have a `train/loss` column if validation'
' metrics computation is successful.')
self.assertLen(
metrics,
total_rounds + 1,
msg='The number of rows in the metrics CSV should be the number of '
'training rounds + 1 (as there is an extra row for validation set'
'metrics after training has completed.')
def test_fedavg_training_decreases_loss(self):
batch = _batch_fn()
federated_data = [[batch]]
iterative_process = _build_federated_averaging_process()
def client_datasets_fn(round_num):
del round_num
return federated_data
def validation_fn(model):
keras_model = tff.simulation.models.mnist.create_keras_model(
compile_model=True)
model.assign_weights_to(keras_model)
return {'loss': keras_model.evaluate(batch.x, batch.y)}
initial_state = iterative_process.initialize()
root_output_dir = self.get_temp_dir()
final_state = training_loop.run(
iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=1,
experiment_name='fedavg_decreases_loss',
root_output_dir=root_output_dir)
self.assertLess(
validation_fn(final_state.model)['loss'],
validation_fn(initial_state.model)['loss'])
def test_raises_non_callable_client_dataset(self):
iterative_process = _build_federated_averaging_process()
client_dataset = [[_batch_fn()]]
def validation_fn(model, round_num):
del model, round_num
return {}
root_output_dir = self.get_temp_dir()
with self.assertRaises(TypeError):
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_dataset,
validation_fn=validation_fn,
total_rounds=10,
experiment_name='non_callable_client_dataset',
root_output_dir=root_output_dir)
def test_raises_non_callable_evaluate_fn(self):
iterative_process = _build_federated_averaging_process()
federated_data = [[_batch_fn()]]
def client_datasets_fn(round_num):
del round_num
return federated_data
metrics_dict = {}
root_output_dir = self.get_temp_dir()
with self.assertRaises(TypeError):
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=metrics_dict,
total_rounds=10,
experiment_name='non_callable_evaluate',
root_output_dir=root_output_dir)
def test_raises_non_str_output_dir(self):
iterative_process = _build_federated_averaging_process()
federated_data = [[_batch_fn()]]
def client_datasets_fn(round_num):
del round_num
return federated_data
def validation_fn(model, round_num):
del model, round_num
return {}
with self.assertRaises(TypeError):
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=10,
experiment_name='non_str_output_dir',
root_output_dir=1)
def test_raises_non_iterative_process(self):
bad_iterative_process = _build_federated_averaging_process().next
federated_data = [[_batch_fn()]]
def client_datasets_fn(round_num):
del round_num
return federated_data
def validation_fn(model, round_num):
del model, round_num
return {}
root_output_dir = self.get_temp_dir()
with self.assertRaises(TypeError):
training_loop.run(iterative_process=[bad_iterative_process],
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=10,
experiment_name='non_iterative_process',
root_output_dir=root_output_dir)
def test_training_loop_works_on_restart(self):
experiment_name = 'test_metrics'
iterative_process = _build_federated_averaging_process()
batch = _batch_fn()
federated_data = [[batch]]
def client_datasets_fn(round_num):
del round_num
return federated_data
def validation_fn(model, round_num):
del model, round_num
return {}
root_output_dir = self.get_temp_dir()
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=2,
experiment_name=experiment_name,
root_output_dir=root_output_dir)
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=4,
experiment_name=experiment_name,
root_output_dir=root_output_dir)
csv_file = os.path.join(root_output_dir, 'results', experiment_name,
'experiment.metrics.csv')
metrics_manager = tff.simulation.CSVMetricsManager(csv_file)
_, metrics = metrics_manager.get_metrics()
# Test that 4 rounds of metrics are logged (the CSV length adds one for
# the header).
self.assertLen(metrics, 5)
def test_checkpoint_manager_saves_state(self):
experiment_name = 'checkpoint_manager_saves_state'
iterative_process = _build_federated_averaging_process()
federated_data = [[_batch_fn()]]
def client_datasets_fn(round_num):
del round_num
return federated_data
def validation_fn(model, round_num):
del model, round_num
return {}
root_output_dir = self.get_temp_dir()
final_state = training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=validation_fn,
total_rounds=1,
experiment_name=experiment_name,
root_output_dir=root_output_dir)
final_model = iterative_process.get_model_weights(final_state)
ckpt_manager = tff.simulation.FileCheckpointManager(
os.path.join(root_output_dir, 'checkpoints', experiment_name))
restored_state, restored_round = ckpt_manager.load_latest_checkpoint(
final_state)
self.assertEqual(restored_round, 0)
keras_model = tff.simulation.models.mnist.create_keras_model(
compile_model=True)
restored_model = iterative_process.get_model_weights(restored_state)
restored_model.assign_weights_to(keras_model)
restored_loss = keras_model.test_on_batch(federated_data[0][0].x,
federated_data[0][0].y)
final_model.assign_weights_to(keras_model)
final_loss = keras_model.test_on_batch(federated_data[0][0].x,
federated_data[0][0].y)
self.assertEqual(final_loss, restored_loss)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('client')
server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('server')
def iterative_process_builder(
model_fn: Callable[[], tff.learning.Model],
) -> tff.templates.IterativeProcess:
"""Creates an iterative process using a given TFF `model_fn`.
Args:
model_fn: A no-arg function returning a `tff.learning.Model`.
Returns:
A `tff.templates.IterativeProcess`.
"""
logging.info('Trainable weights:')
for weight in model_fn().weights.trainable:
logging.info('name: %s shape: %s', weight.name, weight.shape)
if FLAGS.uniform_weighting:
client_weighting = tff.learning.ClientWeighting.UNIFORM
elif FLAGS.task == 'shakespeare' or FLAGS.task == 'stackoverflow_nwp':
def client_weighting(local_outputs):
return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
else:
client_weighting = None
if FLAGS.noise_multiplier is None:
if FLAGS.uniform_weighting:
aggregation_factory = tff.aggregators.UnweightedMeanFactory()
else:
aggregation_factory = tff.aggregators.MeanFactory()
if FLAGS.clip is not None:
if FLAGS.clip <= 0:
raise ValueError('clip must be positive if clipping is enabled.')
if FLAGS.adaptive_clip_learning_rate is None:
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.')
clip = tff.aggregators.PrivateQuantileEstimationProcess.no_noise(
initial_estimate=FLAGS.clip,
target_quantile=FLAGS.target_unclipped_quantile,
learning_rate=FLAGS.adaptive_clip_learning_rate)
aggregation_factory = tff.aggregators.clipping_factory(
clip, aggregation_factory)
else:
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 FLAGS.adaptive_clip_learning_rate is None:
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)
return 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)
task_spec = training_specs.TaskSpec(
iterative_process_builder=iterative_process_builder,
client_epochs_per_round=FLAGS.client_epochs_per_round,
client_batch_size=FLAGS.client_batch_size,
clients_per_round=FLAGS.clients_per_round,
client_datasets_random_seed=FLAGS.client_datasets_random_seed)
if FLAGS.task == 'cifar100':
runner_spec = federated_cifar100.configure_training(task_spec)
elif FLAGS.task == 'emnist_cr':
runner_spec = federated_emnist.configure_training(task_spec)
elif FLAGS.task == 'emnist_ae':
runner_spec = federated_emnist_ae.configure_training(task_spec)
elif FLAGS.task == 'shakespeare':
runner_spec = federated_shakespeare.configure_training(task_spec)
elif FLAGS.task == 'stackoverflow_nwp':
runner_spec = federated_stackoverflow.configure_training(task_spec)
elif FLAGS.task == 'stackoverflow_lr':
runner_spec = federated_stackoverflow_lr.configure_training(task_spec)
else:
raise ValueError(
'--task flag {} is not supported, must be one of {}.'.format(
FLAGS.task, _SUPPORTED_TASKS))
_write_hparam_flags()
training_loop.run(
iterative_process=runner_spec.iterative_process,
client_datasets_fn=runner_spec.client_datasets_fn,
validation_fn=runner_spec.validation_fn,
test_fn=runner_spec.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_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,
vocab_tokens_size: Optional[int] = 10000,
vocab_tags_size: Optional[int] = 500,
max_elements_per_user: Optional[int] = 1000,
num_validation_examples: Optional[int] = 10000,
total_rounds: Optional[int] = 1500,
experiment_name: Optional[str] = 'federated_so_lr',
root_output_dir: Optional[str] = '/tmp/fed_opt',
max_eval_batches: Optional[int] = None,
**kwargs):
"""Runs an iterative process on the Stack Overflow logistic regression task.
This method will load and pre-process dataset and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process that it applies to the task, using
`federated_research.utils.training_loop`.
We assume that the iterative process has the following functional type
signatures:
* `initialize`: `( -> S@SERVER)` where `S` represents the server state.
* `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`
represents the server state, `{B*}` represents the client datasets,
and `T` represents a python `Mapping` object.
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.
vocab_tokens_size: Integer dictating the number of most frequent words to
use in the vocabulary.
vocab_tags_size: Integer dictating the number of most frequent tags to use
in the label creation.
max_elements_per_user: The maximum number of elements processed for each
client's dataset.
num_validation_examples: The number of test examples to use for validation.
total_rounds: The number of federated training rounds.
experiment_name: The name of the experiment being run. This will be appended
to the `root_output_dir` for purposes of writing outputs.
root_output_dir: The name of the root output directory for writing
experiment outputs.
max_eval_batches: If set to a positive integer, evaluation datasets are
capped to at most that many batches. If set to None or a nonpositive
integer, the full evaluation datasets are used.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`federated_research/utils/training_utils.py`.
"""
stackoverflow_train, _ = stackoverflow_tag_prediction.get_federated_datasets(
word_vocab_size=vocab_tokens_size,
tag_vocab_size=vocab_tags_size,
train_client_batch_size=client_batch_size,
train_client_epochs_per_round=client_epochs_per_round,
max_elements_per_train_client=max_elements_per_user)
_, stackoverflow_validation, stackoverflow_test = stackoverflow_tag_prediction.get_centralized_datasets(
train_batch_size=client_batch_size,
word_vocab_size=vocab_tokens_size,
tag_vocab_size=vocab_tags_size,
num_validation_examples=num_validation_examples)
if max_eval_batches and max_eval_batches >= 1:
stackoverflow_validation = stackoverflow_validation.take(max_eval_batches)
stackoverflow_test = stackoverflow_test.take(max_eval_batches)
input_spec = stackoverflow_train.create_tf_dataset_for_client(
stackoverflow_train.client_ids[0]).element_spec
model_builder = functools.partial(
stackoverflow_lr_models.create_logistic_model,
vocab_tokens_size=vocab_tokens_size,
vocab_tags_size=vocab_tags_size)
loss_builder = functools.partial(
tf.keras.losses.BinaryCrossentropy,
from_logits=False,
reduction=tf.keras.losses.Reduction.SUM)
def tff_model_fn() -> tff.learning.Model:
return tff.learning.from_keras_model(
keras_model=model_builder(),
input_spec=input_spec,
loss=loss_builder(),
metrics=metrics_builder())
training_process = iterative_process_builder(tff_model_fn)
client_datasets_fn = training_utils.build_client_datasets_fn(
dataset=stackoverflow_train,
clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed)
evaluate_fn = training_utils.build_centralized_evaluate_fn(
model_builder=model_builder,
eval_dataset=stackoverflow_validation,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
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=stackoverflow_validation.concatenate(stackoverflow_test),
loss_builder=loss_builder,
metrics_builder=metrics_builder)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(
iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
def run_federated(
iterative_process_builder: Callable[...,
tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
schedule: Optional[str] = 'none',
beta: Optional[float] = 0.,
max_batches_per_client: Optional[int] = -1,
client_datasets_random_seed: Optional[int] = None,
model: Optional[str] = 'cnn',
total_rounds: Optional[int] = 1500,
experiment_name: Optional[str] = 'federated_synthetic',
root_output_dir: Optional[str] = '/tmp/fed_opt',
max_eval_batches: Optional[int] = None,
alpha: Optional[float] = 0.,
beta_data: Optional[float] = 0.,
iid: Optional[int] = 0,
num_users: Optional[int] = 100,
**kwargs):
"""Runs an iterative process on the EMNIST character recognition task.
This method will load and pre-process dataset and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process that it applies to the task, using
`federated_research.utils.training_loop`.
We assume that the iterative process has the following functional type
signatures:
* `initialize`: `( -> S@SERVER)` where `S` represents the server state.
* `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`
represents the server state, `{B*}` represents the client datasets,
and `T` represents a python `Mapping` object.
Moreover, the server state must have an attribute `model` of type
`tff.learning.ModelWeights`.
Args:
iterative_process_builder: A function that accepts a no-arg `model_fn`, and
returns a `tff.templates.IterativeProcess`. The `model_fn` must return a
`tff.learning.Model`.
client_epochs_per_round: An integer representing the number of epochs of
training performed per client in each training round.
client_batch_size: An integer representing the batch size used on clients.
clients_per_round: An integer representing the number of clients
participating in each round.
max_batches_per_client: An optional int specifying the number of batches
taken by each client at each round. If `-1`, the entire client dataset is
used.
client_datasets_random_seed: An optional int used to seed which clients are
sampled at each round. If `None`, no seed is used.
model: A string specifying the model used for character recognition.
Can be one of `cnn` and `2nn`, corresponding to a CNN model and a densely
connected 2-layer model (respectively).
total_rounds: The number of federated training rounds.
experiment_name: The name of the experiment being run. This will be appended
to the `root_output_dir` for purposes of writing outputs.
root_output_dir: The name of the root output directory for writing
experiment outputs.
max_eval_batches: If set to a positive integer, evaluation datasets are
capped to at most that many batches. If set to None or a nonpositive
integer, the full evaluation datasets are used.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`federated_research/utils/training_utils.py`.
"""
logging.info(f' DATA PARAMS: ')
logging.info(f' Num Users: {num_users}')
logging.info(f' alpha: {alpha}')
logging.info(f' beta: {beta_data}')
logging.info(f' iid: {iid}')
train_data, test_data, federated_test = synthetic_dataset.generate_federated_softmax_data(
batch_size=client_batch_size,
client_epochs_per_round=client_epochs_per_round,
test_batch_size=100,
alpha=alpha,
beta=beta_data,
iid=iid,
num_users=num_users)
input_spec = train_data.create_tf_dataset_for_client(
train_data.client_ids[0]).element_spec
model_builder = functools.partial(create_logistic_regression_model)
loss_builder = tf.keras.losses.SparseCategoricalCrossentropy
metrics_builder = lambda: [tf.keras.metrics.SparseCategoricalAccuracy()]
def tff_model_fn() -> tff.learning.Model:
return tff.learning.from_keras_model(keras_model=model_builder(),
input_spec=input_spec,
loss=loss_builder(),
metrics=metrics_builder())
training_process = iterative_process_builder(tff_model_fn)
evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=test_data,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
test_fn = training_utils.build_unweighted_test_fn(
federated_eval_dataset=federated_test,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
logging.info('Training model:')
logging.info(model_builder().summary())
try:
var = kwargs['hparam_dict']['var_q_clients']
print(f'Variance: {var}')
q_client = np.load(
f'/home/monica/AVAIL_VECTORS/q_client_{var}_synthetic.npy')
#if q_client is None:
#logging.info('Could not load q_client - initializing random availabilities')
#q_client=None
except:
logging.info(
'Could not load q_client - initializing random availabilities')
q_client = None
if schedule == 'none':
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset=train_data,
train_clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed,
min_clients=kwargs['hparam_dict']['min_clients'],
var_q_clients=kwargs['hparam_dict']['var_q_clients'],
f_mult=kwargs['hparam_dict']['f_mult'],
f_intercept=kwargs['hparam_dict']['f_intercept'],
sine_wave=kwargs['hparam_dict']['sine_wave'],
use_p=True,
q_client=q_client,
)
training_loop.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
elif schedule == 'loss':
if 'loss_pool_size' in kwargs['hparam_dict'] and kwargs['hparam_dict'][
'loss_pool_size'] is not None:
loss_pool_size = kwargs['hparam_dict']['loss_pool_size']
logging.info(f'Loss pool size: {loss_pool_size}')
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset=train_data,
train_clients_per_round=loss_pool_size,
random_seed=client_datasets_random_seed,
min_clients=kwargs['hparam_dict']['min_clients'],
var_q_clients=kwargs['hparam_dict']['var_q_clients'],
f_mult=kwargs['hparam_dict']['f_mult'],
f_intercept=kwargs['hparam_dict']['f_intercept'],
sine_wave=kwargs['hparam_dict']['sine_wave'],
use_p=True,
q_client=q_client)
training_loop_loss.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
total_rounds=total_rounds,
total_clients=loss_pool_size,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
else:
raise ValueError('Loss pool size not specified')
else:
init_p = kwargs['hparam_dict']['initialize_p']
logging.info(f'Initializing as p = {init_p}')
client_datasets_fn = training_utils.build_availability_client_datasets_fn(
train_dataset=train_data,
train_clients_per_round=clients_per_round,
beta=beta,
min_clients=kwargs['hparam_dict']['min_clients'],
var_q_clients=kwargs['hparam_dict']['var_q_clients'],
f_mult=kwargs['hparam_dict']['f_mult'],
f_intercept=kwargs['hparam_dict']['f_intercept'],
sine_wave=kwargs['hparam_dict']['sine_wave'],
q_client=q_client,
initialize_p=init_p,
)
training_loop_importance.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
def run_federated(
iterative_process_builder: Callable[...,
tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
client_datasets_random_seed: Optional[int] = None,
vocab_size: Optional[int] = 10000,
num_oov_buckets: Optional[int] = 1,
sequence_length: Optional[int] = 20,
max_elements_per_user: Optional[int] = 1000,
num_validation_examples: Optional[int] = 10000,
embedding_size: Optional[int] = 96,
latent_size: Optional[int] = 670,
num_layers: Optional[int] = 1,
shared_embedding: Optional[bool] = False,
total_rounds: Optional[int] = 1500,
experiment_name: Optional[str] = 'federated_so_nwp',
root_output_dir: Optional[str] = '/tmp/fed_opt',
**kwargs):
"""Runs an iterative process on the Stack Overflow next word prediction task.
This method will load and pre-process dataset and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process that it applies to the task, using
`federated_research.utils.training_loop`.
We assume that the iterative process has the following functional type
signatures:
* `initialize`: `( -> S@SERVER)` where `S` represents the server state.
* `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`
represents the server state, `{B*}` represents the client datasets,
and `T` represents a python `Mapping` object.
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`, a
`client_weight_fn` and returns a `tff.templates.IterativeProcess`. The
`model_fn` must return a `tff.learning.Model`.
client_epochs_per_round: An integer representing the number of epochs of
training performed per client in each training round.
client_batch_size: An integer representing the batch size used on clients.
clients_per_round: An integer representing the number of clients
participating in each round.
client_datasets_random_seed: An optional int used to seed which clients are
sampled at each round. If `None`, no seed is used.
vocab_size: Integer dictating the number of most frequent words to use in
the vocabulary.
num_oov_buckets: The number of out-of-vocabulary buckets to use.
sequence_length: The maximum number of words to take for each sequence.
max_elements_per_user: The maximum number of elements processed for each
client's dataset.
num_validation_examples: The number of test examples to use for validation.
embedding_size: The dimension of the word embedding layer.
latent_size: The dimension of the latent units in the recurrent layers.
num_layers: The number of stacked recurrent layers to use.
shared_embedding: Boolean indicating whether to tie input and output
embeddings.
total_rounds: The number of federated training rounds.
experiment_name: The name of the experiment being run. This will be appended
to the `root_output_dir` for purposes of writing outputs.
root_output_dir: The name of the root output directory for writing
experiment outputs.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`federated_research/utils/training_utils.py`.
"""
model_builder = functools.partial(
stackoverflow_models.create_recurrent_model,
vocab_size=vocab_size,
num_oov_buckets=num_oov_buckets,
embedding_size=embedding_size,
latent_size=latent_size,
num_layers=num_layers,
shared_embedding=shared_embedding)
loss_builder = functools.partial(
tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True)
special_tokens = stackoverflow_word_prediction.get_special_tokens(
vocab_size, num_oov_buckets)
pad_token = special_tokens.pad
oov_tokens = special_tokens.oov
eos_token = special_tokens.eos
def metrics_builder():
return [
keras_metrics.MaskedCategoricalAccuracy(name='accuracy_with_oov',
masked_tokens=[pad_token]),
keras_metrics.MaskedCategoricalAccuracy(name='accuracy_no_oov',
masked_tokens=[pad_token] +
oov_tokens),
# Notice BOS never appears in ground truth.
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_no_oov_or_eos',
masked_tokens=[pad_token, eos_token] + oov_tokens),
keras_metrics.NumBatchesCounter(),
keras_metrics.NumTokensCounter(masked_tokens=[pad_token])
]
train_clientdata, _, _ = tff.simulation.datasets.stackoverflow.load_data()
# TODO(b/161914546): consider moving evaluation to use
# `tff.learning.build_federated_evaluation` to get metrics over client
# distributions, as well as the example weight means from this centralized
# evaluation.
_, validation_dataset, test_dataset = stackoverflow_word_prediction.get_centralized_datasets(
vocab_size=vocab_size,
max_sequence_length=sequence_length,
num_validation_examples=num_validation_examples,
num_oov_buckets=num_oov_buckets)
train_dataset_preprocess_comp = stackoverflow_word_prediction.create_preprocess_fn(
vocab=stackoverflow_word_prediction.create_vocab(vocab_size),
num_oov_buckets=num_oov_buckets,
client_batch_size=client_batch_size,
client_epochs_per_round=client_epochs_per_round,
max_sequence_length=sequence_length,
max_elements_per_client=max_elements_per_user)
input_spec = train_dataset_preprocess_comp.type_signature.result.element
def tff_model_fn() -> tff.learning.Model:
return tff.learning.from_keras_model(keras_model=model_builder(),
input_spec=input_spec,
loss=loss_builder(),
metrics=metrics_builder())
def client_weight_fn(local_outputs):
# Num_tokens is a tensor with type int64[1], to use as a weight need
# a float32 scalar.
return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
iterative_process = iterative_process_builder(
tff_model_fn, client_weight_fn=client_weight_fn)
training_process = tff.simulation.compose_dataset_computation_with_iterative_process(
train_dataset_preprocess_comp, iterative_process)
training_process.get_model_weights = iterative_process.get_model_weights
client_datasets_fn = training_utils.build_client_datasets_fn(
dataset=train_clientdata,
clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed)
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 model_weights, round_num: evaluate_fn(model_weights)
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)
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 run_federated(
iterative_process_builder: Callable[...,
tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
max_batches_per_client: Optional[int] = -1,
client_datasets_random_seed: Optional[int] = None,
sequence_length: Optional[int] = 80,
total_rounds: Optional[int] = 1500,
experiment_name: Optional[str] = 'federated_shakespeare',
root_output_dir: Optional[str] = '/tmp/fed_opt',
max_eval_batches: Optional[int] = None,
**kwargs):
"""Runs an iterative process on a Shakespeare next character prediction task.
This method will load and pre-process dataset and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process that it applies to the task, using
`federated_research.utils.training_loop`.
We assume that the iterative process has the following functional type
signatures:
* `initialize`: `( -> S@SERVER)` where `S` represents the server state.
* `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`
represents the server state, `{B*}` represents the client datasets,
and `T` represents a python `Mapping` object.
Moreover, the server state must have an attribute `model` of type
`tff.learning.ModelWeights`.
Args:
iterative_process_builder: A function that accepts a no-arg `model_fn`, and
a `client_weight_fn`, and returns a `tff.templates.IterativeProcess`. The
`model_fn` must return a `tff.learning.Model`.
client_epochs_per_round: An integer representing the number of epochs of
training performed per client in each training round.
client_batch_size: An integer representing the batch size used on clients.
clients_per_round: An integer representing the number of clients
participating in each round.
max_batches_per_client: An optional int specifying the number of batches
taken by each client at each round. If `-1`, the entire client dataset is
used.
client_datasets_random_seed: An optional int used to seed which clients are
sampled at each round. If `None`, no seed is used.
sequence_length: An int specifying the length of the character sequences
used for prediction.
total_rounds: The number of federated training rounds.
experiment_name: The name of the experiment being run. This will be appended
to the `root_output_dir` for purposes of writing outputs.
root_output_dir: The name of the root output directory for writing
experiment outputs.
max_eval_batches: If set to a positive integer, evaluation datasets are
capped to at most that many batches. If set to None or a nonpositive
integer, the full evaluation datasets are used.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`federated_research/utils/training_utils.py`.
"""
train_clientdata = shakespeare_dataset.construct_character_level_datasets(
client_batch_size=client_batch_size,
client_epochs_per_round=client_epochs_per_round,
sequence_length=sequence_length,
max_batches_per_client=max_batches_per_client)
_, test_dataset = shakespeare_dataset.get_centralized_datasets(
train_batch_size=client_batch_size,
max_test_batches=max_eval_batches,
sequence_length=sequence_length)
model_builder = functools.partial(create_shakespeare_model,
sequence_length=sequence_length)
loss_builder = functools.partial(
tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True)
input_spec = train_clientdata.create_tf_dataset_for_client(
train_clientdata.client_ids[0]).element_spec
def client_weight_fn(local_outputs):
# Num_tokens is a tensor with type int64[1], to use as a weight need
# a float32 scalar.
return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
def tff_model_fn() -> tff.learning.Model:
return tff.learning.from_keras_model(keras_model=model_builder(),
input_spec=input_spec,
loss=loss_builder(),
metrics=metrics_builder())
training_process = iterative_process_builder(
tff_model_fn, client_weight_fn=client_weight_fn)
client_datasets_fn = training_utils.build_client_datasets_fn(
dataset=train_clientdata,
clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed)
evaluate_fn = training_utils.build_centralized_evaluate_fn(
eval_dataset=test_dataset,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=evaluate_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
def run_experiment():
"""Data preprocessing and experiment execution."""
emnist_train, emnist_test = emnist_dataset.get_emnist_datasets(
FLAGS.client_batch_size,
FLAGS.client_epochs_per_round,
only_digits=FLAGS.only_digits)
example_dataset = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0])
input_spec = example_dataset.element_spec
client_datasets_fn = training_utils.build_client_datasets_fn(
emnist_train, FLAGS.clients_per_round)
evaluate_fn = training_utils.build_centralized_evaluate_fn(
eval_dataset=emnist_test,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
client_optimizer_fn = functools.partial(
utils_impl.create_optimizer_from_flags, 'client')
server_optimizer_fn = functools.partial(
utils_impl.create_optimizer_from_flags, 'server')
def tff_model_fn():
keras_model = model_builder()
return tff.learning.from_keras_model(keras_model,
input_spec=input_spec,
loss=loss_builder(),
metrics=metrics_builder())
if FLAGS.use_compression:
# We create a `MeasuredProcess` for broadcast process and a
# `MeasuredProcess` for aggregate process by providing the
# `_broadcast_encoder_fn` and `_mean_encoder_fn` to corresponding utilities.
# The fns are called once for each of the model weights created by
# tff_model_fn, and return instances of appropriate encoders.
encoded_broadcast_process = (
tff.learning.framework.build_encoded_broadcast_process_from_model(
tff_model_fn, _broadcast_encoder_fn))
encoded_mean_process = (
tff.learning.framework.build_encoded_mean_process_from_model(
tff_model_fn, _mean_encoder_fn))
else:
encoded_broadcast_process = None
encoded_mean_process = None
iterative_process = tff.learning.build_federated_averaging_process(
model_fn=tff_model_fn,
client_optimizer_fn=client_optimizer_fn,
server_optimizer_fn=server_optimizer_fn,
aggregation_process=encoded_mean_process,
broadcast_process=encoded_broadcast_process)
hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags())
training_loop_dict = utils_impl.lookup_flag_values(training_loop_flags)
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
hparam_dict=hparam_dict,
**training_loop_dict)
def run_federated(
iterative_process_builder: Callable[..., tff.templates.IterativeProcess],
evaluation_computation_builder: Callable[..., tff.Computation],
*, # Caller passes below args by name.
client_batch_size: int,
clients_per_round: int,
global_variables_only: bool,
# Begin task-specific parameters.
split_by_user: bool,
split_train_fraction: float,
split_val_fraction: float,
normalize_ratings: bool,
max_examples_per_user: Optional[int],
num_items: int,
num_latent_factors: int,
add_biases: bool,
l2_regularization: float,
spreadout_lambda: float,
accuracy_threshold: float,
dataset_path:
str = 'http://files.grouplens.org/datasets/movielens/ml-1m.zip',
# End task-specific parameters.
total_rounds: int,
experiment_name: Optional[str] = 'federated_ml_mf',
root_output_dir: Optional[str] = '/tmp/fed_recon',
**kwargs):
"""Runs an iterative process on the MovieLens matrix factorization 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`.
This algorithm only sends updates for item embeddings every round. User
embeddings are reconstructed every round based on the latest item embeddings
and user data.
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`, a
`loss_fn`, a `metrics_fn`, and a `client_weight_fn`, and returns a
`tff.templates.IterativeProcess`. The `model_fn` must return a
`reconstruction_model.ReconstructionModel`. See `federated_trainer.py` for
an example.
evaluation_computation_builder: A function that accepts a no-arg `model_fn`,
a loss_fn`, and a `metrics_fn`, and returns a `tff.Computation` for
federated reconstruction evaluation. The `model_fn` must return a
`reconstruction_model.ReconstructionModel`. See `federated_trainer.py` for
an example.
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.
global_variables_only: If True, the `ReconstructionModel` contains
all model variables as global variables. This can be useful for
baselines involving aggregating all variables.
split_by_user: Whether to split MovieLens data into train/val/test by user
ID or by timestamp. If True, `movielens_dataset.split_tf_datasets` is used
to partition the set of users into disjoint train/val/test sets. If False,
`movielens_dataset.split_ratings_df` is used to split each user's data
into train/val/test portions, so that each user shows up in each data
partition. Setting to False can be useful for comparing with server-side
training, since server matrix factorization requires test users to have
been seen before (since otherwise we don't have trained user embeddings
for these users).
split_train_fraction: The fraction of data to use for the train set.
split_val_fraction: The fraction of data to use for the val set.
`split_train_fraction` and `split_val_fraction` should sum to no more than
1. 1 - split_train_fraction - split_val_fraction of the data is left for
test.
normalize_ratings: Whether to normalize ratings in 1-5 to be in {-1, -0.5,
0, 0.5, 1} via a linear scaling.
max_examples_per_user: If not None, limit the number of rating examples for
each user to this many examples.
num_items: Number of items in the preferences matrix.
num_latent_factors: Dimensionality of the learned user/item embeddings used
to factorize the preferences matrix.
add_biases: If True, add three bias terms: (1) user-specific bias, (2)
item-specific bias, and (3) global bias.
l2_regularization: The constant to use to scale L2 regularization on all
weights, including the factorized matrices and the (optional) biases. A
value of 0.0 indicates no regularization.
spreadout_lambda: Scaling constant for spreadout regularization on item
embeddings. This ensures that item embeddings are generally spread far
apart, and that random items have dissimilar embeddings. See
`models.EmbeddingSpreadoutRegularizer` for details. A value of 0.0
indicates no regularization.
accuracy_threshold: Threshold to use to determine whether a prediction is
considered correct for metrics.
dataset_path: URL or local path to the MovieLens 1M dataset. If a URL is
passed, it is expected to be a .zip archive that will be extracted.
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 training_loop.py`.
"""
logging.info('Copying MovieLens data.')
if tf.io.gfile.exists('/tmp/ml-1m'):
tf.io.gfile.rmtree('/tmp/ml-1m')
if dataset_path.startswith('http'):
movielens_dataset.download_and_extract_data(dataset_path, '/tmp')
else:
tf.io.gfile.makedirs('/tmp/ml-1m/')
tf.io.gfile.copy(
os.path.join(dataset_path, 'ratings.dat'),
'/tmp/ml-1m/ratings.dat',
overwrite=True)
tf.io.gfile.copy(
os.path.join(dataset_path, 'movies.dat'),
'/tmp/ml-1m/movies.dat',
overwrite=True)
tf.io.gfile.copy(
os.path.join(dataset_path, 'users.dat'),
'/tmp/ml-1m/users.dat',
overwrite=True)
logging.info('Finished copying MovieLens data.')
ratings_df, _, _ = movielens_dataset.load_movielens_data(
normalize_ratings=normalize_ratings)
# Split the ratings into training/val/test.
if split_by_user:
tf_datasets = movielens_dataset.create_tf_datasets(
ratings_df=ratings_df,
personal_model=True,
batch_size=client_batch_size,
max_examples_per_user=max_examples_per_user,
num_local_epochs=1)
tf_train_datasets, tf_val_datasets, tf_test_datasets = movielens_dataset.split_tf_datasets(
tf_datasets,
train_fraction=split_train_fraction,
val_fraction=split_val_fraction)
else:
train_ratings_df, val_ratings_df, test_ratings_df = movielens_dataset.split_ratings_df(
ratings_df,
train_fraction=split_train_fraction,
val_fraction=split_val_fraction)
tf_train_datasets = movielens_dataset.create_tf_datasets(
ratings_df=train_ratings_df,
personal_model=True,
batch_size=client_batch_size,
max_examples_per_user=max_examples_per_user,
num_local_epochs=1)
tf_val_datasets = movielens_dataset.create_tf_datasets(
ratings_df=val_ratings_df,
personal_model=True,
batch_size=client_batch_size,
max_examples_per_user=max_examples_per_user,
num_local_epochs=1)
tf_test_datasets = movielens_dataset.create_tf_datasets(
ratings_df=test_ratings_df,
personal_model=True,
batch_size=client_batch_size,
max_examples_per_user=max_examples_per_user,
num_local_epochs=1)
model_fn = models.build_reconstruction_model(
functools.partial(
models.get_matrix_factorization_model,
num_users=1,
num_items=num_items,
num_latent_factors=num_latent_factors,
personal_model=True,
add_biases=add_biases,
l2_regularization=l2_regularization,
spreadout_lambda=spreadout_lambda),
global_variables_only=global_variables_only)
loss_fn = models.get_loss_fn()
metrics_fn = models.get_metrics_fn(accuracy_threshold=accuracy_threshold)
training_process = iterative_process_builder(
model_fn, loss_fn=loss_fn, metrics_fn=metrics_fn)
evaluation_computation = evaluation_computation_builder(
model_fn, loss_fn=loss_fn, metrics_fn=metrics_fn)
def client_datasets_fn_from_tf_datasets(
tf_datasets: List[tf.data.Dataset],
clients_per_round: int,
) -> Callable[[int], List[tf.data.Dataset]]:
"""Produces a sampling function for train/val/test from a list of datasets."""
sample_clients_fn = federated_trainer_utils.build_list_sample_fn(
list(range(len(tf_datasets))), size=clients_per_round, replace=False)
def client_datasets_fn(round_num):
sampled_clients = sample_clients_fn(round_num)
return [tf_datasets[client_id] for client_id in sampled_clients]
return client_datasets_fn
# Create client sampling functions for each of train/val/test.
train_client_datasets_fn = client_datasets_fn_from_tf_datasets(
tf_train_datasets, clients_per_round=clients_per_round)
val_client_datasets_fn = client_datasets_fn_from_tf_datasets(
tf_val_datasets, clients_per_round=clients_per_round)
test_client_datasets_fn = client_datasets_fn_from_tf_datasets(
tf_test_datasets, clients_per_round=clients_per_round)
# Create final evaluation functions to pass to `training_loop`.
val_fn = federated_trainer_utils.build_eval_fn(
evaluation_computation=evaluation_computation,
client_datasets_fn=val_client_datasets_fn)
test_fn = federated_trainer_utils.build_eval_fn(
evaluation_computation=evaluation_computation,
client_datasets_fn=test_client_datasets_fn)
test_fn = functools.partial(test_fn, round_num=0)
logging.info('Starting training loop.')
training_loop.run(
iterative_process=training_process,
client_datasets_fn=train_client_datasets_fn,
validation_fn=val_fn,
test_fn=test_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
def run_federated(
iterative_process_builder: Callable[..., tff.templates.IterativeProcess],
evaluation_computation_builder: Callable[..., tff.Computation],
client_batch_size: int,
clients_per_round: int,
global_variables_only: bool,
vocab_size: int = 10000,
num_oov_buckets: int = 1,
sequence_length: int = 20,
max_elements_per_user: int = 1000,
embedding_size: int = 96,
latent_size: int = 670,
num_layers: int = 1,
total_rounds: int = 1500,
experiment_name: str = 'federated_so_nwp',
root_output_dir: str = '/tmp/fed_recon',
split_dataset_strategy: str = federated_trainer_utils
.SPLIT_STRATEGY_AGGREGATED,
split_dataset_proportion: int = 2,
compose_dataset_computation: bool = False,
**kwargs):
"""Runs an iterative process on the Stack Overflow next word prediction task.
This method will load and pre-process dataset and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process that it applies to the task, using
`federated_research.utils.training_loop`.
This model only sends updates for its embeddings corresponding to the most
common words. Embeddings for out of vocabulary buckets are reconstructed on
device at the beginning of each round, and destroyed at the end of these
rounds.
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`, a
`loss_fn`, a `metrics_fn`, and a `client_weight_fn`, and returns a
`tff.templates.IterativeProcess`. The `model_fn` must return a
`reconstruction_model.ReconstructionModel`. See `federated_trainer.py` for
an example.
evaluation_computation_builder: A function that accepts a no-arg `model_fn`,
a loss_fn`, and a `metrics_fn`, and returns a `tff.Computation` for
federated reconstruction evaluation. The `model_fn` must return a
`reconstruction_model.ReconstructionModel`. See `federated_trainer.py` for
an example.
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.
global_variables_only: If True, the `ReconstructionModel` contains all model
variables as global variables. This can be useful for baselines involving
aggregating all variables.
vocab_size: Integer dictating the number of most frequent words to use in
the vocabulary.
num_oov_buckets: The number of out-of-vocabulary buckets to use.
sequence_length: The maximum number of words to take for each sequence.
max_elements_per_user: The maximum number of elements processed for each
client's dataset.
embedding_size: The dimension of the word embedding layer.
latent_size: The dimension of the latent units in the recurrent layers.
num_layers: The number of stacked recurrent layers to use.
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.
split_dataset_strategy: The method to use to split the data. Must be one of
`skip`, in which case every `split_dataset_proportion` example is used for
reconstruction, or `aggregated`, when the first
1/`split_dataset_proportion` proportion of the examples is used for
reconstruction.
split_dataset_proportion: Parameter controlling how much of the data is used
for reconstruction. If `split_dataset_proportion` is n, then 1 / n of the
data is used for reconstruction.
compose_dataset_computation: Whether to compose dataset computation with
training and evaluation computations. If True, may speed up experiments by
parallelizing dataset computations in multimachine setups. Not currently
supported in OSS.
**kwargs: Additional arguments configuring the training loop. For details on
supported arguments, see `training_loop.py`.
"""
loss_fn = functools.partial(
tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True)
special_tokens = stackoverflow_word_prediction.get_special_tokens(
vocab_size, num_oov_buckets)
pad_token = special_tokens.pad
oov_tokens = special_tokens.oov
eos_token = special_tokens.eos
def metrics_fn():
return [
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_with_oov', masked_tokens=[pad_token]),
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_no_oov', masked_tokens=[pad_token] + oov_tokens),
# Notice BOS never appears in ground truth.
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_no_oov_or_eos',
masked_tokens=[pad_token, eos_token] + oov_tokens),
keras_metrics.NumBatchesCounter(),
keras_metrics.NumTokensCounter(masked_tokens=[pad_token])
]
train_clientdata, validation_clientdata, test_clientdata = (
tff.simulation.datasets.stackoverflow.load_data())
vocab = stackoverflow_word_prediction.create_vocab(vocab_size)
dataset_preprocess_comp = stackoverflow_dataset.create_preprocess_fn(
vocab=vocab,
num_oov_buckets=num_oov_buckets,
client_batch_size=client_batch_size,
max_sequence_length=sequence_length,
max_elements_per_client=max_elements_per_user,
feature_dtypes=train_clientdata.element_type_structure,
sort_by_date=True)
input_spec = dataset_preprocess_comp.type_signature.result.element
model_fn = functools.partial(
models.create_recurrent_reconstruction_model,
vocab_size=vocab_size,
num_oov_buckets=num_oov_buckets,
embedding_size=embedding_size,
latent_size=latent_size,
num_layers=num_layers,
input_spec=input_spec,
global_variables_only=global_variables_only)
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)
iterative_process = iterative_process_builder(
model_fn,
loss_fn=loss_fn,
metrics_fn=metrics_fn,
client_weight_fn=client_weight_fn,
dataset_split_fn_builder=functools.partial(
federated_trainer_utils.build_dataset_split_fn,
split_dataset_strategy=split_dataset_strategy,
split_dataset_proportion=split_dataset_proportion))
base_eval_computation = evaluation_computation_builder(
model_fn,
loss_fn=loss_fn,
metrics_fn=metrics_fn,
dataset_split_fn_builder=functools.partial(
federated_trainer_utils.build_dataset_split_fn,
split_dataset_strategy=split_dataset_strategy,
split_dataset_proportion=split_dataset_proportion))
if compose_dataset_computation:
# Compose dataset computations with client training and evaluation to avoid
# linear cost of computing centrally. This changes the expected input of
# the `IterativeProcess` and `tff.Computation` to be a list of client IDs
# instead of datasets.
training_process = (
tff.simulation.compose_dataset_computation_with_iterative_process(
dataset_preprocess_comp, iterative_process))
training_process = (
tff.simulation.compose_dataset_computation_with_iterative_process(
train_clientdata.dataset_computation, training_process))
training_process.get_model_weights = iterative_process.get_model_weights
base_eval_computation = (
tff.simulation.compose_dataset_computation_with_computation(
dataset_preprocess_comp, base_eval_computation))
val_computation = (
tff.simulation.compose_dataset_computation_with_computation(
validation_clientdata.dataset_computation, base_eval_computation))
test_computation = (
tff.simulation.compose_dataset_computation_with_computation(
test_clientdata.dataset_computation, base_eval_computation))
# Create client sampling functions for each of train/val/test.
# We need to sample client IDs, not datasets, and we do not need to apply
# `dataset_preprocess_comp` since this is applied as part of the training
# process and evaluation computation.
train_client_datasets_fn = federated_trainer_utils.build_list_sample_fn(
train_clientdata.client_ids, size=clients_per_round, replace=False)
val_client_datasets_fn = federated_trainer_utils.build_list_sample_fn(
validation_clientdata.client_ids, size=clients_per_round, replace=False)
test_client_datasets_fn = federated_trainer_utils.build_list_sample_fn(
test_clientdata.client_ids, size=clients_per_round, replace=False)
else:
training_process = iterative_process
val_computation = base_eval_computation
test_computation = base_eval_computation
# Apply dataset computations.
train_clientdata = train_clientdata.preprocess(dataset_preprocess_comp)
validation_clientdata = validation_clientdata.preprocess(
dataset_preprocess_comp)
test_clientdata = test_clientdata.preprocess(dataset_preprocess_comp)
# Create client sampling functions for each of train/val/test.
train_client_datasets_fn = functools.partial(
tff.simulation.build_uniform_sampling_fn(train_clientdata.client_ids),
size=clients_per_round)
val_client_datasets_fn = functools.partial(
tff.simulation.build_uniform_sampling_fn(
validation_clientdata.client_ids),
size=clients_per_round)
test_client_datasets_fn = functools.partial(
tff.simulation.build_uniform_sampling_fn(test_clientdata.client_ids),
size=clients_per_round)
# Create final evaluation functions to pass to `training_loop`.
val_fn = federated_trainer_utils.build_eval_fn(
evaluation_computation=val_computation,
client_datasets_fn=val_client_datasets_fn,
get_model=training_process.get_model_weights)
test_fn = federated_trainer_utils.build_eval_fn(
evaluation_computation=test_computation,
client_datasets_fn=test_client_datasets_fn,
get_model=training_process.get_model_weights)
test_fn = functools.partial(test_fn, round_num=0)
training_loop.run(
iterative_process=training_process,
client_datasets_fn=train_client_datasets_fn,
validation_fn=val_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=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))
tff.backends.native.set_local_execution_context(max_fanout=10)
model_builder = functools.partial(
stackoverflow_models.create_recurrent_model,
vocab_size=FLAGS.vocab_size,
embedding_size=FLAGS.embedding_size,
latent_size=FLAGS.latent_size,
num_layers=FLAGS.num_layers,
shared_embedding=FLAGS.shared_embedding)
loss_builder = functools.partial(
tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True)
special_tokens = stackoverflow_dataset.get_special_tokens(FLAGS.vocab_size)
pad_token = special_tokens.pad
oov_tokens = special_tokens.oov
eos_token = special_tokens.eos
def metrics_builder():
return [
keras_metrics.MaskedCategoricalAccuracy(name='accuracy_with_oov',
masked_tokens=[pad_token]),
keras_metrics.MaskedCategoricalAccuracy(name='accuracy_no_oov',
masked_tokens=[pad_token] +
oov_tokens),
# Notice BOS never appears in ground truth.
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_no_oov_or_eos',
masked_tokens=[pad_token, eos_token] + oov_tokens),
keras_metrics.NumBatchesCounter(),
keras_metrics.NumTokensCounter(masked_tokens=[pad_token]),
]
datasets = stackoverflow_dataset.construct_word_level_datasets(
FLAGS.vocab_size, FLAGS.client_batch_size,
FLAGS.client_epochs_per_round, FLAGS.sequence_length,
FLAGS.max_elements_per_user, FLAGS.num_validation_examples)
train_dataset, validation_dataset, test_dataset = datasets
if FLAGS.uniform_weighting:
def client_weight_fn(local_outputs):
del local_outputs
return 1.0
else:
def client_weight_fn(local_outputs):
return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
def model_fn():
return tff.learning.from_keras_model(
model_builder(),
loss_builder(),
input_spec=validation_dataset.element_spec,
metrics=metrics_builder())
if FLAGS.noise_multiplier is not None:
if not FLAGS.uniform_weighting:
raise ValueError(
'Differential privacy is only implemented for uniform weighting.'
)
dp_query = tff.utils.build_dp_query(
clip=FLAGS.clip,
noise_multiplier=FLAGS.noise_multiplier,
expected_total_weight=FLAGS.clients_per_round,
adaptive_clip_learning_rate=FLAGS.adaptive_clip_learning_rate,
target_unclipped_quantile=FLAGS.target_unclipped_quantile,
clipped_count_budget_allocation=FLAGS.
clipped_count_budget_allocation,
expected_clients_per_round=FLAGS.clients_per_round)
weights_type = tff.learning.framework.weights_type_from_model(model_fn)
aggregation_process = tff.utils.build_dp_aggregate_process(
weights_type.trainable, dp_query)
else:
aggregation_process = None
server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'server')
client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'client')
iterative_process = tff.learning.build_federated_averaging_process(
model_fn=model_fn,
server_optimizer_fn=server_optimizer_fn,
client_weight_fn=client_weight_fn,
client_optimizer_fn=client_optimizer_fn,
aggregation_process=aggregation_process)
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset, FLAGS.clients_per_round)
evaluate_fn = training_utils.build_centralized_evaluate_fn(
model_builder=model_builder,
eval_dataset=validation_dataset,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
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)
logging.info('Training model:')
logging.info(model_builder().summary())
hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags())
training_loop_dict = utils_impl.lookup_flag_values(training_loop_flags)
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
hparam_dict=hparam_dict,
**training_loop_dict)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'client')
server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'server')
client_lr_callback = callbacks.create_reduce_lr_on_plateau(
learning_rate=FLAGS.client_learning_rate,
decay_factor=FLAGS.client_decay_factor,
min_delta=FLAGS.min_delta,
min_lr=FLAGS.min_lr,
window_size=FLAGS.window_size,
patience=FLAGS.patience)
server_lr_callback = callbacks.create_reduce_lr_on_plateau(
learning_rate=FLAGS.server_learning_rate,
decay_factor=FLAGS.server_decay_factor,
min_delta=FLAGS.min_delta,
min_lr=FLAGS.min_lr,
window_size=FLAGS.window_size,
patience=FLAGS.patience)
def iterative_process_builder(
model_fn: Callable[[], tff.learning.Model],
) -> tff.templates.IterativeProcess:
"""Creates an iterative process using a given TFF `model_fn`.
Args:
model_fn: A no-arg function returning a `tff.learning.Model`.
Returns:
A `tff.templates.IterativeProcess`.
"""
return adaptive_fed_avg.build_fed_avg_process(
model_fn,
client_lr_callback,
server_lr_callback,
client_optimizer_fn=client_optimizer_fn,
server_optimizer_fn=server_optimizer_fn)
task_spec = training_specs.TaskSpec(
iterative_process_builder=iterative_process_builder,
client_epochs_per_round=FLAGS.client_epochs_per_round,
client_batch_size=FLAGS.client_batch_size,
clients_per_round=FLAGS.clients_per_round,
client_datasets_random_seed=FLAGS.client_datasets_random_seed)
if FLAGS.task == 'cifar100':
runner_spec = federated_cifar100.configure_training(
task_spec, crop_size=FLAGS.cifar100_crop_size)
elif FLAGS.task == 'emnist_cr':
runner_spec = federated_emnist.configure_training(
task_spec, model=FLAGS.emnist_cr_model)
elif FLAGS.task == 'emnist_ae':
runner_spec = federated_emnist_ae.configure_training(task_spec)
elif FLAGS.task == 'shakespeare':
runner_spec = federated_shakespeare.configure_training(
task_spec, sequence_length=FLAGS.shakespeare_sequence_length)
elif FLAGS.task == 'stackoverflow_nwp':
runner_spec = federated_stackoverflow.configure_training(
task_spec,
vocab_size=FLAGS.so_nwp_vocab_size,
num_oov_buckets=FLAGS.so_nwp_num_oov_buckets,
sequence_length=FLAGS.so_nwp_sequence_length,
max_elements_per_user=FLAGS.so_nwp_max_elements_per_user,
num_validation_examples=FLAGS.so_nwp_num_validation_examples)
elif FLAGS.task == 'stackoverflow_lr':
runner_spec = federated_stackoverflow_lr.configure_training(
task_spec,
vocab_tokens_size=FLAGS.so_lr_vocab_tokens_size,
vocab_tags_size=FLAGS.so_lr_vocab_tags_size,
max_elements_per_user=FLAGS.so_lr_max_elements_per_user,
num_validation_examples=FLAGS.so_lr_num_validation_examples)
else:
raise ValueError(
'--task flag {} is not supported, must be one of {}.'.format(
FLAGS.task, _SUPPORTED_TASKS))
training_loop.run(iterative_process=runner_spec.iterative_process,
client_datasets_fn=runner_spec.client_datasets_fn,
validation_fn=runner_spec.validation_fn,
test_fn=runner_spec.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,
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))
client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'client')
server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'server')
client_lr_schedule = optimizer_utils.create_lr_schedule_from_flags(
'client')
server_lr_schedule = optimizer_utils.create_lr_schedule_from_flags(
'server')
def iterative_process_builder(
model_fn: Callable[[], tff.learning.Model],
client_weight_fn: Optional[Callable[[Any], tf.Tensor]] = None,
) -> tff.templates.IterativeProcess:
"""Creates an iterative process using a given TFF `model_fn`.
Args:
model_fn: A no-arg function returning a `tff.learning.Model`.
client_weight_fn: Optional function that takes the output of
`model.report_local_outputs` and returns a tensor providing the weight
in the federated average of model deltas. If not provided, the default
is the total number of examples processed on device.
Returns:
A `tff.templates.IterativeProcess`.
"""
return fed_avg_schedule.build_fed_avg_process(
model_fn=model_fn,
client_optimizer_fn=client_optimizer_fn,
client_lr=client_lr_schedule,
server_optimizer_fn=server_optimizer_fn,
server_lr=server_lr_schedule,
client_weight_fn=client_weight_fn)
task_spec = training_specs.TaskSpec(
iterative_process_builder=iterative_process_builder,
client_epochs_per_round=FLAGS.client_epochs_per_round,
client_batch_size=FLAGS.client_batch_size,
clients_per_round=FLAGS.clients_per_round,
client_datasets_random_seed=FLAGS.client_datasets_random_seed)
if FLAGS.task == 'cifar100':
runner_spec = federated_cifar100.configure_training(
task_spec,
crop_size=FLAGS.cifar100_crop_size,
distort_train_images=FLAGS.cifar100_distort_train_images)
elif FLAGS.task == 'emnist_cr':
runner_spec = federated_emnist.configure_training(
task_spec, model=FLAGS.emnist_cr_model)
elif FLAGS.task == 'emnist_ae':
runner_spec = federated_emnist_ae.configure_training(task_spec)
elif FLAGS.task == 'shakespeare':
runner_spec = federated_shakespeare.configure_training(
task_spec, sequence_length=FLAGS.shakespeare_sequence_length)
elif FLAGS.task == 'stackoverflow_nwp':
runner_spec = federated_stackoverflow.configure_training(
task_spec,
vocab_size=FLAGS.so_nwp_vocab_size,
num_oov_buckets=FLAGS.so_nwp_num_oov_buckets,
sequence_length=FLAGS.so_nwp_sequence_length,
max_elements_per_user=FLAGS.so_nwp_max_elements_per_user,
num_validation_examples=FLAGS.so_nwp_num_validation_examples)
elif FLAGS.task == 'stackoverflow_lr':
runner_spec = federated_stackoverflow_lr.configure_training(
task_spec,
vocab_tokens_size=FLAGS.so_lr_vocab_tokens_size,
vocab_tags_size=FLAGS.so_lr_vocab_tags_size,
max_elements_per_user=FLAGS.so_lr_max_elements_per_user,
num_validation_examples=FLAGS.so_lr_num_validation_examples)
else:
raise ValueError(
'--task flag {} is not supported, must be one of {}.'.format(
FLAGS.task, _SUPPORTED_TASKS))
_write_hparam_flags()
training_loop.run(iterative_process=runner_spec.iterative_process,
client_datasets_fn=runner_spec.client_datasets_fn,
validation_fn=runner_spec.validation_fn,
test_fn=runner_spec.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,
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_test = emnist_dataset.get_emnist_datasets(
FLAGS.client_batch_size,
FLAGS.client_epochs_per_round,
only_digits=False)
if FLAGS.model == 'cnn':
model_builder = functools.partial(
emnist_models.create_conv_dropout_model, only_digits=False)
elif FLAGS.model == '2nn':
model_builder = functools.partial(
emnist_models.create_two_hidden_layer_model, only_digits=False)
else:
raise ValueError('Cannot handle model flag [{!s}].'.format(
FLAGS.model))
loss_builder = tf.keras.losses.SparseCategoricalCrossentropy
metrics_builder = lambda: [tf.keras.metrics.SparseCategoricalAccuracy()]
if FLAGS.uniform_weighting:
def client_weight_fn(local_outputs):
del local_outputs
return 1.0
else:
client_weight_fn = None # Defaults to the number of examples per client.
def model_fn():
return tff.learning.from_keras_model(
model_builder(),
loss_builder(),
input_spec=emnist_test.element_spec,
metrics=metrics_builder())
if FLAGS.noise_multiplier is not None:
if not FLAGS.uniform_weighting:
raise ValueError(
'Differential privacy is only implemented for uniform weighting.'
)
dp_query = tff.utils.build_dp_query(
clip=FLAGS.clip,
noise_multiplier=FLAGS.noise_multiplier,
expected_total_weight=FLAGS.clients_per_round,
adaptive_clip_learning_rate=FLAGS.adaptive_clip_learning_rate,
target_unclipped_quantile=FLAGS.target_unclipped_quantile,
clipped_count_budget_allocation=FLAGS.
clipped_count_budget_allocation,
expected_clients_per_round=FLAGS.clients_per_round,
per_vector_clipping=FLAGS.per_vector_clipping,
model=model_fn())
weights_type = tff.learning.framework.weights_type_from_model(model_fn)
aggregation_process = tff.utils.build_dp_aggregate_process(
weights_type.trainable, dp_query)
else:
aggregation_process = None
server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'server')
client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'client')
iterative_process = tff.learning.build_federated_averaging_process(
model_fn=model_fn,
server_optimizer_fn=server_optimizer_fn,
client_weight_fn=client_weight_fn,
client_optimizer_fn=client_optimizer_fn,
aggregation_process=aggregation_process)
client_datasets_fn = training_utils.build_client_datasets_fn(
emnist_train, FLAGS.clients_per_round)
evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=emnist_test,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
logging.info('Training model:')
logging.info(model_builder().summary())
hparam_dict = utils_impl.lookup_flag_values(utils_impl.get_hparam_flags())
training_loop_dict = utils_impl.lookup_flag_values(training_loop_flags)
training_loop.run(iterative_process=iterative_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
hparam_dict=hparam_dict,
**training_loop_dict)
def run_federated(
iterative_process_builder: Callable[...,
tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
max_batches_per_client: Optional[int] = -1,
client_datasets_random_seed: Optional[int] = None,
total_rounds: Optional[int] = 1500,
experiment_name: Optional[str] = 'federated_emnist_ae',
root_output_dir: Optional[str] = '/tmp/fed_opt',
max_eval_batches: Optional[int] = None,
**kwargs):
"""Runs an iterative process on the EMNIST autoencoder task.
This method will load and pre-process dataset and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process that it applies to the task, using
`federated_research.utils.training_loop`.
We assume that the iterative process has the following functional type
signatures:
* `initialize`: `( -> S@SERVER)` where `S` represents the server state.
* `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`
represents the server state, `{B*}` represents the client datasets,
and `T` represents a python `Mapping` object.
Moreover, the server state must have an attribute `model` of type
`tff.learning.ModelWeights`.
Args:
iterative_process_builder: A function that accepts a no-arg `model_fn`, and
returns a `tff.templates.IterativeProcess`. The `model_fn` must return a
`tff.learning.Model`.
client_epochs_per_round: An integer representing the number of epochs of
training performed per client in each training round.
client_batch_size: An integer representing the batch size used on clients.
clients_per_round: An integer representing the number of clients
participating in each round.
max_batches_per_client: An optional int specifying the number of batches
taken by each client at each round. If `-1`, the entire client dataset is
used.
client_datasets_random_seed: An optional int used to seed which clients are
sampled at each round. If `None`, no seed is used.
total_rounds: The number of federated training rounds.
experiment_name: The name of the experiment being run. This will be appended
to the `root_output_dir` for purposes of writing outputs.
root_output_dir: The name of the root output directory for writing
experiment outputs.
max_eval_batches: If set to a positive integer, evaluation datasets are
capped to at most that many batches. If set to None or a nonpositive
integer, the full evaluation datasets are used.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`federated_research/utils/training_utils.py`.
"""
emnist_train, _ = emnist_ae_dataset.get_emnist_datasets(
client_batch_size=client_batch_size,
client_epochs_per_round=client_epochs_per_round,
max_batches_per_client=max_batches_per_client,
only_digits=False)
_, emnist_test = emnist_ae_dataset.get_centralized_datasets(
train_batch_size=client_batch_size,
max_test_batches=max_eval_batches,
only_digits=False)
input_spec = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0]).element_spec
model_builder = emnist_ae_models.create_autoencoder_model
loss_builder = functools.partial(tf.keras.losses.MeanSquaredError,
reduction=tf.keras.losses.Reduction.SUM)
metrics_builder = lambda: [tf.keras.metrics.MeanSquaredError()]
def tff_model_fn() -> tff.learning.Model:
return tff.learning.from_keras_model(keras_model=model_builder(),
input_spec=input_spec,
loss=loss_builder(),
metrics=metrics_builder())
training_process = iterative_process_builder(tff_model_fn)
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset=emnist_train,
train_clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed)
evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=emnist_test,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=evaluate_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
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_federated(
iterative_process_builder: Callable[...,
tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
schedule: Optional[str] = 'none',
beta: Optional[float] = 0.,
max_batches_per_client: Optional[int] = -1,
client_datasets_random_seed: Optional[int] = None,
crop_size: Optional[int] = 24,
total_rounds: Optional[int] = 1500,
experiment_name: Optional[str] = 'federated_cifar100',
root_output_dir: Optional[str] = '/tmp/fed_opt',
max_eval_batches: Optional[int] = None,
**kwargs):
"""Runs an iterative process on the CIFAR-100 classification task.
This method will load and pre-process dataset and construct a model used for
the task. It then uses `iterative_process_builder` to create an iterative
process that it applies to the task, using
`federated_research.utils.training_loop`.
We assume that the iterative process has the following functional type
signatures:
* `initialize`: `( -> S@SERVER)` where `S` represents the server state.
* `next`: `<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>` where `S`
represents the server state, `{B*}` represents the client datasets,
and `T` represents a python `Mapping` object.
Moreover, the server state must have an attribute `model` of type
`tff.learning.ModelWeights`.
Args:
iterative_process_builder: A function that accepts a no-arg `model_fn`, and
returns a `tff.templates.IterativeProcess`. The `model_fn` must return a
`tff.learning.Model`.
client_epochs_per_round: An integer representing the number of epochs of
training performed per client in each training round.
client_batch_size: An integer representing the batch size used on clients.
clients_per_round: An integer representing the number of clients
participating in each round.
max_batches_per_client: An optional int specifying the number of batches
taken by each client at each round. If `-1`, the entire client dataset is
used.
client_datasets_random_seed: An optional int used to seed which clients are
sampled at each round. If `None`, no seed is used.
crop_size: An optional integer representing the resulting size of input
images after preprocessing.
total_rounds: The number of federated training rounds.
experiment_name: The name of the experiment being run. This will be appended
to the `root_output_dir` for purposes of writing outputs.
root_output_dir: The name of the root output directory for writing
experiment outputs.
max_eval_batches: If set to a positive integer, evaluation datasets are
capped to at most that many batches. If set to None or a nonpositive
integer, the full evaluation datasets are used.
**kwargs: Additional arguments configuring the training loop. For details
on supported arguments, see
`federated_research/utils/training_utils.py`.
"""
crop_shape = (crop_size, crop_size, 3)
cifar_train, _, fed_test_data = cifar100_dataset.get_federated_cifar100(
client_epochs_per_round=client_epochs_per_round,
train_batch_size=client_batch_size,
crop_shape=crop_shape,
max_batches_per_client=max_batches_per_client)
_, cifar_test = cifar100_dataset.get_centralized_datasets(
train_batch_size=client_batch_size,
max_test_batches=max_eval_batches,
crop_shape=crop_shape)
input_spec = cifar_train.create_tf_dataset_for_client(
cifar_train.client_ids[0]).element_spec
model_builder = functools.partial(resnet_models.create_resnet18,
input_shape=crop_shape,
num_classes=NUM_CLASSES)
loss_builder = tf.keras.losses.SparseCategoricalCrossentropy
metrics_builder = lambda: [tf.keras.metrics.SparseCategoricalAccuracy()]
def tff_model_fn() -> tff.learning.Model:
return tff.learning.from_keras_model(keras_model=model_builder(),
input_spec=input_spec,
loss=loss_builder(),
metrics=metrics_builder())
training_process = iterative_process_builder(tff_model_fn)
evaluate_fn = training_utils.build_evaluate_fn(
eval_dataset=cifar_test,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
test_fn = training_utils.build_unweighted_test_fn(
federated_eval_dataset=fed_test_data,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder)
logging.info('Training model:')
logging.info(model_builder().summary())
try:
var = kwargs['hparam_dict']['var_q_clients']
q_client = np.load(
f'/home/monica/AVAIL_VECTORS/q_client_{var}_cifar.npy')
except:
logging.info(
'Could not load q_client - initializing random availabilities')
q_client = None
if schedule == 'none':
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset=cifar_train,
train_clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed,
var_q_clients=kwargs['hparam_dict']['var_q_clients'],
f_mult=kwargs['hparam_dict']['f_mult'],
f_intercept=kwargs['hparam_dict']['f_intercept'],
sine_wave=kwargs['hparam_dict']['sine_wave'],
use_p=True,
q_client=q_client,
)
training_loop.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
elif schedule == 'loss':
if 'loss_pool_size' in kwargs['hparam_dict'] and kwargs['hparam_dict'][
'loss_pool_size'] is not None:
loss_pool_size = kwargs['hparam_dict']['loss_pool_size']
logging.info(f'Loss pool size: {loss_pool_size}')
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset=cifar_train,
train_clients_per_round=loss_pool_size,
random_seed=client_datasets_random_seed,
var_q_clients=kwargs['hparam_dict']['var_q_clients'],
f_mult=kwargs['hparam_dict']['f_mult'],
f_intercept=kwargs['hparam_dict']['f_intercept'],
sine_wave=kwargs['hparam_dict']['sine_wave'],
use_p=True,
q_client=q_client,
)
training_loop_loss.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
total_rounds=total_rounds,
total_clients=loss_pool_size,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
else:
raise ValueError('Loss pool size not specified')
else:
client_datasets_fn = training_utils.build_availability_client_datasets_fn(
train_dataset=cifar_train,
train_clients_per_round=clients_per_round,
random_seed=client_datasets_random_seed,
beta=beta,
var_q_clients=kwargs['hparam_dict']['var_q_clients'],
f_mult=kwargs['hparam_dict']['f_mult'],
f_intercept=kwargs['hparam_dict']['f_intercept'],
sine_wave=kwargs['hparam_dict']['sine_wave'],
q_client=q_client,
)
training_loop_importance.run(iterative_process=training_process,
client_datasets_fn=client_datasets_fn,
validation_fn=evaluate_fn,
test_fn=test_fn,
total_rounds=total_rounds,
experiment_name=experiment_name,
root_output_dir=root_output_dir,
**kwargs)
def run_federated(
iterative_process_builder: Callable[..., tff.templates.IterativeProcess],
client_epochs_per_round: int,
client_batch_size: int,
clients_per_round: int,
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)
