def test_oov_token_correct(self):
vocab = ['A', 'B', 'C']
max_seq_len = 5
to_ids_fn = dataset.build_to_ids_fn(vocab, max_seq_len)
_, oov_token, _, _ = dataset.get_special_tokens(len(vocab))
data = {'tokens': 'A B D'}
processed = to_ids_fn(data)
self.assertEqual(self.evaluate(processed)[3], oov_token)
def test_build_to_ids_fn_truncates(self):
vocab = ['A', 'B', 'C']
max_seq_len = 1
_, _, bos, _ = dataset.get_special_tokens(len(vocab))
to_ids_fn = dataset.build_to_ids_fn(vocab, max_seq_len)
data = {'tokens': 'A B C'}
processed = to_ids_fn(data)
self.assertAllEqual(self.evaluate(processed), [bos, 1])
def test_build_to_ids_fn_embeds_all_vocab(self):
vocab = ['A', 'B', 'C']
max_seq_len = 5
_, _, bos, eos = dataset.get_special_tokens(len(vocab))
to_ids_fn = dataset.build_to_ids_fn(vocab, max_seq_len)
data = {'tokens': 'A B C'}
processed = to_ids_fn(data)
self.assertAllEqual(self.evaluate(processed), [bos, 1, 2, 3, eos])
def test_pad_token_correct(self):
vocab = ['A', 'B', 'C']
max_seq_len = 5
to_ids_fn = dataset.build_to_ids_fn(vocab, max_seq_len)
pad, _, bos, eos = dataset.get_special_tokens(len(vocab))
data = {'tokens': 'A B C'}
processed = to_ids_fn(data)
batched_ds = tf.data.Dataset.from_tensor_slices([processed]).padded_batch(
1, padded_shapes=[6])
sample_elem = next(iter(batched_ds))
self.assertAllEqual(self.evaluate(sample_elem), [[bos, 1, 2, 3, eos, pad]])
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
tf.compat.v1.enable_v2_behavior()
model_builder = functools.partial(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)
pad_token, oov_token, _, eos_token = dataset.get_special_tokens(
FLAGS.vocab_size)
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_token]),
# Notice BOS never appears in ground truth.
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_no_oov_or_eos',
masked_tokens=[pad_token, oov_token, eos_token]),
keras_metrics.NumBatchesCounter(),
keras_metrics.NumTokensCounter(masked_tokens=[pad_token])
]
train_set, validation_set, test_set = 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)
input_spec = validation_set.element_spec
def client_weight_fn(local_outputs):
# Num_tokens is a tensor with type int64[1], to use as a weight need
# a float32 scalar.
return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
training_process = iterative_process_builder.from_flags(
input_spec=input_spec,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder,
client_weight_fn=client_weight_fn)
client_datasets_fn = training_utils.build_client_datasets_fn(
train_set, FLAGS.clients_per_round)
assign_weights_fn = fed_avg_schedule.ServerState.assign_weights_to_keras_model
evaluate_fn = training_utils.build_evaluate_fn(
model_builder=model_builder,
eval_dataset=validation_set,
loss_builder=loss_builder,
metrics_builder=metrics_builder,
assign_weights_to_keras_model=assign_weights_fn)
test_fn = training_utils.build_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_set.concatenate(test_set),
loss_builder=loss_builder,
metrics_builder=metrics_builder,
assign_weights_to_keras_model=assign_weights_fn)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(training_process,
client_datasets_fn,
evaluate_fn,
test_fn=test_fn)
def run_experiment():
"""Runs the training experiment."""
(_, stackoverflow_validation,
stackoverflow_test) = dataset.construct_word_level_datasets(
FLAGS.vocab_size, FLAGS.batch_size, 1, FLAGS.sequence_length, -1,
FLAGS.num_validation_examples)
centralized_train = dataset.get_centralized_train_dataset(
FLAGS.vocab_size, FLAGS.batch_size, FLAGS.sequence_length,
FLAGS.shuffle_buffer_size)
def _lstm_fn(latent_size):
return tf.keras.layers.LSTM(latent_size, return_sequences=True)
model = models.create_recurrent_model(
FLAGS.vocab_size,
_lstm_fn,
'stackoverflow-lstm',
shared_embedding=FLAGS.shared_embedding)
logging.info('Training model: %s', model.summary())
optimizer = optimizer_utils.create_optimizer_fn_from_flags('centralized')()
pad_token, oov_token, _, eos_token = dataset.get_special_tokens(
FLAGS.vocab_size)
model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=optimizer,
metrics=[
# Plus 4 for pad, oov, bos, eos
keras_metrics.FlattenedCategoricalAccuracy(
vocab_size=FLAGS.vocab_size + 4,
name='accuracy_with_oov',
masked_tokens=pad_token),
keras_metrics.FlattenedCategoricalAccuracy(
vocab_size=FLAGS.vocab_size + 4,
name='accuracy_no_oov',
masked_tokens=[pad_token, oov_token]),
keras_metrics.FlattenedCategoricalAccuracy(
vocab_size=FLAGS.vocab_size + 4,
name='accuracy_no_oov_or_eos',
masked_tokens=[pad_token, oov_token, eos_token]),
])
train_results_path = os.path.join(FLAGS.root_output_dir, 'train_results',
FLAGS.experiment_name)
test_results_path = os.path.join(FLAGS.root_output_dir, 'test_results',
FLAGS.experiment_name)
train_csv_logger = keras_callbacks.AtomicCSVLogger(train_results_path)
test_csv_logger = keras_callbacks.AtomicCSVLogger(test_results_path)
log_dir = os.path.join(FLAGS.root_output_dir, 'logdir',
FLAGS.experiment_name)
try:
tf.io.gfile.makedirs(log_dir)
tf.io.gfile.makedirs(train_results_path)
tf.io.gfile.makedirs(test_results_path)
except tf.errors.OpError:
pass # log_dir already exists.
train_tensorboard_callback = tf.keras.callbacks.TensorBoard(
log_dir=log_dir,
write_graph=True,
update_freq=FLAGS.tensorboard_update_frequency)
test_tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir)
# Write the hyperparameters to a CSV:
hparam_dict = collections.OrderedDict([(name, FLAGS[name].value)
for name in hparam_flags])
hparams_file = os.path.join(FLAGS.root_output_dir, FLAGS.experiment_name,
'hparams.csv')
utils_impl.atomic_write_to_csv(pd.Series(hparam_dict), hparams_file)
model.fit(centralized_train,
epochs=FLAGS.epochs,
verbose=0,
validation_data=stackoverflow_validation,
callbacks=[train_csv_logger, train_tensorboard_callback])
score = model.evaluate(
stackoverflow_test,
verbose=0,
callbacks=[test_csv_logger, test_tensorboard_callback])
logging.info('Final test loss: %.4f', score[0])
logging.info('Final test accuracy: %.4f', score[1])
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
tf.compat.v1.enable_v2_behavior()
tff.framework.set_default_executor(
tff.framework.local_executor_factory(max_fanout=10))
model_builder = functools.partial(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)
pad_token, oov_token, _, eos_token = dataset.get_special_tokens(
FLAGS.vocab_size)
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_token]),
# Notice BOS never appears in ground truth.
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_no_oov_or_eos',
masked_tokens=[pad_token, oov_token, eos_token]),
keras_metrics.NumBatchesCounter(),
keras_metrics.NumTokensCounter(masked_tokens=[pad_token]),
]
datasets = 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)
if FLAGS.noise_multiplier is not None:
if not FLAGS.uniform_weighting:
raise ValueError(
'Differential privacy is only implemented for uniform weighting.'
)
dp_query = tff.utils.build_dp_query(
clip=FLAGS.clip,
noise_multiplier=FLAGS.noise_multiplier,
expected_total_weight=FLAGS.clients_per_round,
adaptive_clip_learning_rate=FLAGS.adaptive_clip_learning_rate,
target_unclipped_quantile=FLAGS.target_unclipped_quantile,
clipped_count_budget_allocation=FLAGS.
clipped_count_budget_allocation,
expected_num_clients=FLAGS.clients_per_round,
per_vector_clipping=FLAGS.per_vector_clipping,
model=model_builder())
dp_aggregate_fn, _ = tff.utils.build_dp_aggregate(dp_query)
else:
dp_aggregate_fn = None
def model_fn():
return tff.learning.from_keras_model(
model_builder(),
loss_builder(),
input_spec=validation_dataset.element_spec,
metrics=metrics_builder())
server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'server')
client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags(
'client')
training_process = (
tff.learning.federated_averaging.build_federated_averaging_process(
model_fn=model_fn,
server_optimizer_fn=server_optimizer_fn,
client_weight_fn=client_weight_fn,
client_optimizer_fn=client_optimizer_fn,
stateful_delta_aggregate_fn=dp_aggregate_fn))
adaptive_clipping = (FLAGS.adaptive_clip_learning_rate > 0)
training_process = DPFedAvgProcessAdapter(training_process,
FLAGS.per_vector_clipping,
adaptive_clipping)
client_datasets_fn = training_utils.build_client_datasets_fn(
train_dataset, FLAGS.clients_per_round)
evaluate_fn = training_utils.build_evaluate_fn(
model_builder=model_builder,
eval_dataset=validation_dataset,
loss_builder=loss_builder,
metrics_builder=metrics_builder,
assign_weights_to_keras_model=assign_weights_to_keras_model)
test_fn = training_utils.build_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,
assign_weights_to_keras_model=assign_weights_to_keras_model)
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(training_process,
client_datasets_fn,
evaluate_fn,
test_fn=test_fn)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Expected no command-line arguments, '
'got: {}'.format(argv))
tf.compat.v1.enable_v2_behavior()
tff.framework.set_default_executor(
tff.framework.local_executor_factory(max_fanout=10))
if FLAGS.lstm:
def _layer_fn(x):
return tf.keras.layers.LSTM(x, return_sequences=True)
else:
def _layer_fn(x):
return tf.keras.layers.GRU(x, return_sequences=True)
model_builder = functools.partial(models.create_recurrent_model,
vocab_size=FLAGS.vocab_size,
recurrent_layer_fn=_layer_fn,
shared_embedding=FLAGS.shared_embedding)
loss_builder = functools.partial(
tf.keras.losses.SparseCategoricalCrossentropy, from_logits=True)
pad_token, oov_token, _, eos_token = dataset.get_special_tokens(
FLAGS.vocab_size)
def metrics_builder():
return [
keras_metrics.FlattenedCategoricalAccuracy(
# Plus 4 for PAD, OOV, BOS and EOS.
vocab_size=FLAGS.vocab_size + 4,
name='accuracy_with_oov',
masked_tokens=pad_token),
keras_metrics.FlattenedCategoricalAccuracy(
vocab_size=FLAGS.vocab_size + 4,
name='accuracy_no_oov',
masked_tokens=[pad_token, oov_token]),
# Notice BOS never appears in ground truth.
keras_metrics.FlattenedCategoricalAccuracy(
vocab_size=FLAGS.vocab_size + 4,
name='accuracy_no_oov_or_eos',
masked_tokens=[pad_token, oov_token, eos_token]),
keras_metrics.NumBatchesCounter(),
keras_metrics.FlattenedNumExamplesCounter(name='num_tokens',
mask_zero=True),
]
(stackoverflow_train, stackoverflow_validation,
stackoverflow_test) = 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)
sample_batch = tf.nest.map_structure(lambda x: x.numpy(),
next(iter(stackoverflow_validation)))
def client_weight_fn(local_outputs):
# Num_tokens is a tensor with type int64[1], to use as a weight need
# a float32 scalar.
return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
training_process = iterative_process_builder.from_flags(
dummy_batch=sample_batch,
model_builder=model_builder,
loss_builder=loss_builder,
metrics_builder=metrics_builder,
client_weight_fn=client_weight_fn)
client_datasets_fn = training_utils.build_client_datasets_fn(
stackoverflow_train, FLAGS.clients_per_round)
eval_fn = training_utils.build_evaluate_fn(
model_builder=model_builder,
eval_dataset=stackoverflow_validation,
loss_builder=loss_builder,
metrics_builder=metrics_builder,
# Use both val and test for symmetry with other experiments, which
# evaluate on the entire test set.
test_dataset=stackoverflow_validation.concatenate(stackoverflow_test))
logging.info('Training model:')
logging.info(model_builder().summary())
training_loop.run(training_process, client_datasets_fn, eval_fn)
