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)