def test_preprocess_applied(self, mock_load_data, mock_load_word_counts): if tf.config.list_logical_devices('GPU'): self.skipTest('skip GPU test') # Mock out the actual data loading from disk. Assert that the preprocessing # function is applied to the client data, and that only the ClientData # objects we desired are used. # # The correctness of the preprocessing function is tested in other tests. sample_ds = tf.data.Dataset.from_tensor_slices(TEST_DATA) mock_train = mock.create_autospec(tff.simulation.datasets.ClientData) mock_train.create_tf_dataset_from_all_clients = mock.Mock( return_value=sample_ds) mock_validation = mock.create_autospec( tff.simulation.datasets.ClientData) mock_test = mock.create_autospec(tff.simulation.datasets.ClientData) mock_test.create_tf_dataset_from_all_clients = mock.Mock( return_value=sample_ds) mock_load_data.return_value = (mock_train, mock_validation, mock_test) # Return a factor word dictionary. mock_load_word_counts.return_value = collections.OrderedDict(a=1) _, _, _ = stackoverflow_word_prediction.get_centralized_datasets( vocab_size=1000, train_batch_size=10, validation_batch_size=50, test_batch_size=100, num_validation_examples=10000, max_sequence_length=20, num_oov_buckets=1) # Assert the validation ClientData isn't used. mock_load_data.assert_called_once() self.assertEmpty(mock_validation.mock_calls) # Assert the validation ClientData isn't used, and the train and test # are amalgamated into datasets single datasets over all clients. mock_load_data.assert_called_once() self.assertEmpty(mock_validation.mock_calls) self.assertEqual( mock_train.mock_calls, mock.call.create_tf_dataset_from_all_clients().call_list()) self.assertEqual( mock_test.mock_calls, mock.call.create_tf_dataset_from_all_clients().call_list()) # Assert the word counts were loaded once to apply to each dataset. mock_load_word_counts.assert_called_once()
def run_centralized(optimizer: tf.keras.optimizers.Optimizer, num_epochs: int, batch_size: int, decay_epochs: Optional[int] = None, lr_decay: Optional[float] = None, vocab_size: int = 10000, num_oov_buckets: int = 1, d_embed: int = 96, d_model: int = 512, d_hidden: int = 2048, num_heads: int = 8, num_layers: int = 1, max_position_encoding: int = 1000, dropout: float = 0.1, num_validation_examples: int = 10000, sequence_length: int = 20, experiment_name: str = 'centralized_stackoverflow', root_output_dir: str = '/tmp/fedopt_guide', hparams_dict: Optional[Mapping[str, Any]] = None, max_batches: Optional[int] = None): """Trains an Transformer on the Stack Overflow next word prediction task. Args: optimizer: A `tf.keras.optimizers.Optimizer` used to perform training. num_epochs: The number of training epochs. batch_size: The batch size, used for train, validation, and test. decay_epochs: The number of epochs of training before decaying the learning rate. If None, no decay occurs. lr_decay: The amount to decay the learning rate by after `decay_epochs` training epochs have occurred. vocab_size: Vocab size for normal tokens. num_oov_buckets: Number of out of vocabulary buckets. d_embed: Dimension of the token embeddings. d_model: Dimension of features of MultiHeadAttention layers. d_hidden: Dimension of hidden layers of the FFN. num_heads: Number of attention heads. num_layers: Number of Transformer blocks. max_position_encoding: Maximum number of positions for position embeddings. dropout: Dropout rate. num_validation_examples: The number of test examples to use for validation. sequence_length: The maximum number of words to take for each sequence. experiment_name: The name of the experiment. Part of the output directory. root_output_dir: The top-level output directory for experiment runs. The `experiment_name` argument will be appended, and the directory will contain tensorboard logs, metrics written as CSVs, and a CSV of hyperparameter choices (if `hparams_dict` is used). hparams_dict: A mapping with string keys representing the hyperparameters and their values. If not None, this is written to CSV. max_batches: If set to a positive integer, datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full datasets are used. """ train_dataset, validation_dataset, test_dataset = stackoverflow_word_prediction.get_centralized_datasets( vocab_size, sequence_length, train_batch_size=batch_size, num_validation_examples=num_validation_examples, num_oov_buckets=num_oov_buckets, ) if max_batches and max_batches >= 1: train_dataset = train_dataset.take(max_batches) validation_dataset = validation_dataset.take(max_batches) test_dataset = test_dataset.take(max_batches) model = transformer_models.create_transformer_lm( vocab_size=vocab_size, num_oov_buckets=num_oov_buckets, d_embed=d_embed, d_model=d_model, d_hidden=d_hidden, num_heads=num_heads, num_layers=num_layers, max_position_encoding=max_position_encoding, dropout=dropout, name='stackoverflow-transformer') special_tokens = stackoverflow_word_prediction.get_special_tokens( vocab_size=vocab_size, num_oov_buckets=num_oov_buckets) pad_token = special_tokens.pad oov_tokens = special_tokens.oov eos_token = special_tokens.eos model.compile( loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=optimizer, metrics=[ keras_metrics.MaskedCategoricalAccuracy(name='accuracy_with_oov', masked_tokens=[pad_token]), keras_metrics.MaskedCategoricalAccuracy(name='accuracy_no_oov', masked_tokens=[pad_token] + oov_tokens), keras_metrics.MaskedCategoricalAccuracy( name='accuracy_no_oov_or_eos', masked_tokens=[pad_token, eos_token] + oov_tokens), ]) centralized_training_loop.run(keras_model=model, train_dataset=train_dataset, validation_dataset=validation_dataset, test_dataset=test_dataset, experiment_name=experiment_name, root_output_dir=root_output_dir, num_epochs=num_epochs, hparams_dict=hparams_dict, decay_epochs=decay_epochs, lr_decay=lr_decay)
def run_centralized(optimizer: tf.keras.optimizers.Optimizer, experiment_name: str, root_output_dir: str, num_epochs: int, batch_size: int, decay_epochs: Optional[int] = None, lr_decay: Optional[float] = None, hparams_dict: Optional[Mapping[str, Any]] = None, vocab_size: Optional[int] = 10000, num_oov_buckets: Optional[int] = 1, sequence_length: Optional[int] = 20, 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, max_batches: Optional[int] = None, cache_dir: Optional[str] = None): """Trains an RNN on the Stack Overflow next word prediction task. Args: optimizer: A `tf.keras.optimizers.Optimizer` used to perform training. experiment_name: The name of the experiment. Part of the output directory. root_output_dir: The top-level output directory for experiment runs. The `experiment_name` argument will be appended, and the directory will contain tensorboard logs, metrics written as CSVs, and a CSV of hyperparameter choices (if `hparams_dict` is used). num_epochs: The number of training epochs. batch_size: The batch size, used for train, validation, and test. decay_epochs: The number of epochs of training before decaying the learning rate. If None, no decay occurs. lr_decay: The amount to decay the learning rate by after `decay_epochs` training epochs have occurred. hparams_dict: A mapping with string keys representing the hyperparameters and their values. If not None, this is written to CSV. 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. 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. max_batches: If set to a positive integer, datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full datasets are used. """ train_dataset, validation_dataset, test_dataset = stackoverflow_word_prediction.get_centralized_datasets( vocab_size=vocab_size, max_sequence_length=sequence_length, train_batch_size=batch_size, num_validation_examples=num_validation_examples, num_oov_buckets=num_oov_buckets, cache_dir=cache_dir) if max_batches and max_batches >= 1: train_dataset = train_dataset.take(max_batches) validation_dataset = validation_dataset.take(max_batches) test_dataset = test_dataset.take(max_batches) model = stackoverflow_models.create_recurrent_model( vocab_size=vocab_size, num_oov_buckets=num_oov_buckets, name='stackoverflow-lstm', embedding_size=embedding_size, latent_size=latent_size, num_layers=num_layers, shared_embedding=shared_embedding) special_tokens = stackoverflow_word_prediction.get_special_tokens( vocab_size=vocab_size, num_oov_buckets=num_oov_buckets) pad_token = special_tokens.pad oov_tokens = special_tokens.oov eos_token = special_tokens.eos model.compile( loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=optimizer, metrics=[ keras_metrics.MaskedCategoricalAccuracy(name='accuracy_with_oov', masked_tokens=[pad_token]), keras_metrics.MaskedCategoricalAccuracy(name='accuracy_no_oov', masked_tokens=[pad_token] + oov_tokens), keras_metrics.MaskedCategoricalAccuracy( name='accuracy_no_oov_or_eos', masked_tokens=[pad_token, eos_token] + oov_tokens), ]) centralized_training_loop.run(keras_model=model, train_dataset=train_dataset, validation_dataset=validation_dataset, test_dataset=test_dataset, experiment_name=experiment_name, root_output_dir=root_output_dir, num_epochs=num_epochs, hparams_dict=hparams_dict, decay_epochs=decay_epochs, lr_decay=lr_decay)
def 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 configure_training( task_spec: training_specs.TaskSpec, vocab_size: int = 10000, num_oov_buckets: int = 1, sequence_length: int = 20, max_elements_per_user: int = 1000, num_validation_examples: int = 10000, embedding_size: int = 96, latent_size: int = 670, num_layers: int = 1, shared_embedding: bool = False) -> training_specs.RunnerSpec: """Configures training for Stack Overflow next-word prediction. This method will load and pre-process datasets and construct a model used for the task. It then uses `iterative_process_builder` to create an iterative process compatible with `federated_research.utils.training_loop`. Args: task_spec: A `TaskSpec` class for creating federated training tasks. 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. Returns: A `RunnerSpec` containing attributes used for running the newly created federated task. """ 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=task_spec.client_batch_size, client_epochs_per_round=task_spec.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()) iterative_process = task_spec.iterative_process_builder(tff_model_fn) @tff.tf_computation(tf.string) def train_dataset_computation(client_id): client_train_data = train_clientdata.dataset_computation(client_id) return train_dataset_preprocess_comp(client_train_data) training_process = tff.simulation.compose_dataset_computation_with_iterative_process( train_dataset_computation, iterative_process) client_ids_fn = training_utils.build_sample_fn( train_clientdata.client_ids, size=task_spec.clients_per_round, replace=False, random_seed=task_spec.client_datasets_random_seed) # 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)) training_process.get_model_weights = iterative_process.get_model_weights centralized_validation_fn = training_utils.build_centralized_evaluate_fn( model_builder=model_builder, eval_dataset=validation_dataset, loss_builder=loss_builder, metrics_builder=metrics_builder) def validation_fn(server_state, round_num): del round_num return centralized_validation_fn( iterative_process.get_model_weights(server_state)) centralized_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) def test_fn(server_state): return centralized_test_fn( iterative_process.get_model_weights(server_state)) return training_specs.RunnerSpec(iterative_process=training_process, client_datasets_fn=client_sampling_fn, validation_fn=validation_fn, test_fn=test_fn)
def run_federated(iterative_process_builder: Callable[ ..., tff.templates.IterativeProcess], client_epochs_per_round: int, client_batch_size: int, clients_per_round: int, max_elements_per_user: int, total_rounds: int = 3000, vocab_size: int = 10000, num_oov_buckets: int = 1, sequence_length: int = 20, num_validation_examples: int = 10000, dim_embed: int = 96, dim_model: int = 512, dim_hidden: int = 2048, num_heads: int = 8, num_layers: int = 1, max_position_encoding: int = 1000, dropout: float = 0.1, client_datasets_random_seed: Optional[int] = None, experiment_name: str = 'federated_stackoverflow', root_output_dir: str = '/tmp/fedopt_guide', max_val_test_batches: Optional[int] = None, **kwargs) -> None: """Configures training for Stack Overflow next-word prediction. 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/fedopt_guide/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. max_elements_per_user: The maximum number of elements processed for each client's dataset. This has be to a positive value or -1 (which means that all elements are taken for training). total_rounds: The number of federated training rounds. 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. num_validation_examples: The number of test examples to use for validation. dim_embed: An integer for the dimension of the token embeddings. dim_model: An integer for the dimension of features of MultiHeadAttention layers. dim_hidden: An integer for the dimension of hidden layers of the FFN. num_heads: An integer for the number of attention heads. num_layers: An integer for the number of Transformer blocks. max_position_encoding: Maximum number of positions for position embeddings. dropout: Dropout rate. client_datasets_random_seed: An optional int used to seed which clients are sampled at each round. If `None`, no seed is used. 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_val_test_batches: If set to a positive integer, val and test datasets are capped to at most that many batches. If set to None or a nonpositive integer, the full datasets are used. **kwargs: Additional arguments configuring the training loop. For details on supported arguments, see `federated_research/fedopt_guide/training_utils.py`. Returns: A `RunnerSpec` containing attributes used for running the newly created federated task. """ train_clientdata, _, _ = tff.simulation.datasets.stackoverflow.load_data() _, 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) if max_val_test_batches and max_val_test_batches >= 1: validation_dataset = validation_dataset.take(max_val_test_batches) test_dataset = test_dataset.take(max_val_test_batches) model_builder = functools.partial( transformer_models.create_transformer_lm, vocab_size=vocab_size, num_oov_buckets=num_oov_buckets, dim_embed=dim_embed, dim_model=dim_model, dim_hidden=dim_hidden, num_heads=num_heads, num_layers=num_layers, max_position_encoding=max_position_encoding, dropout=dropout, name='stackoverflow-transformer') 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_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) if hasattr(train_clientdata, 'dataset_computation'): @tff.tf_computation(tf.string) def train_dataset_computation(client_id): client_train_data = train_clientdata.dataset_computation(client_id) return train_dataset_preprocess_comp(client_train_data) training_process = tff.simulation.compose_dataset_computation_with_iterative_process( train_dataset_computation, iterative_process) client_ids_fn = tff.simulation.build_uniform_sampling_fn( train_clientdata.client_ids, size=clients_per_round, replace=False, random_seed=client_datasets_random_seed) # 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)) else: training_process = tff.simulation.compose_dataset_computation_with_iterative_process( train_dataset_preprocess_comp, iterative_process) client_sampling_fn = tff.simulation.build_uniform_client_sampling_fn( dataset=train_clientdata, clients_per_round=clients_per_round, random_seed=client_datasets_random_seed) training_process.get_model_weights = iterative_process.get_model_weights evaluate_fn = tff.learning.build_federated_evaluation(tff_model_fn) def validation_fn(model_weights, round_num): del round_num return evaluate_fn(model_weights, [validation_dataset]) def test_fn(model_weights): return evaluate_fn(model_weights, [validation_dataset.concatenate(test_dataset)]) logging.info('Training model:') logging.info(model_builder().summary()) training_loop.run(iterative_process=training_process, train_client_datasets_fn=client_sampling_fn, evaluation_fn=validation_fn, test_fn=test_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)) 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_word_prediction.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]), ] train_dataset, _ = stackoverflow_word_prediction.get_federated_datasets( vocab_size=FLAGS.vocab_size, train_client_batch_size=FLAGS.client_batch_size, train_client_epochs_per_round=FLAGS.client_epochs_per_round, max_sequence_length=FLAGS.sequence_length, max_elements_per_train_client=FLAGS.max_elements_per_user) _, validation_dataset, test_dataset = stackoverflow_word_prediction.get_centralized_datasets( vocab_size=FLAGS.vocab_size, max_sequence_length=FLAGS.sequence_length, num_validation_examples=FLAGS.num_validation_examples) 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) 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()) 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=validation_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)) 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_word_prediction.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]), ] train_dataset, _ = stackoverflow_word_prediction.get_federated_datasets( vocab_size=FLAGS.vocab_size, train_client_batch_size=FLAGS.client_batch_size, train_client_epochs_per_round=FLAGS.client_epochs_per_round, max_sequence_length=FLAGS.sequence_length, max_elements_per_train_client=FLAGS.max_elements_per_user) _, validation_dataset, test_dataset = stackoverflow_word_prediction.get_centralized_datasets( vocab_size=FLAGS.vocab_size, max_sequence_length=FLAGS.sequence_length, num_validation_examples=FLAGS.num_validation_examples) 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=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.' ) 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( 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) validation_fn = lambda state, round_num: evaluate_fn(state.model) evaluate_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) test_fn = lambda state: evaluate_test_fn(state.model) 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, 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)