def test_raises_length_2_crop(self): with self.assertRaises(ValueError): cifar100_dataset.get_federated_cifar100(client_epochs_per_round=1, train_batch_size=10, crop_shape=(32, 32)) with self.assertRaises(ValueError): cifar100_dataset.get_centralized_datasets(train_batch_size=10, crop_shape=(32, 32))
def test_preprocess_applied(self, mock_load_data): 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.ClientData) mock_train.create_tf_dataset_from_all_clients = mock.Mock( return_value=sample_ds) mock_test = mock.create_autospec(tff.simulation.ClientData) mock_test.create_tf_dataset_from_all_clients = mock.Mock( return_value=sample_ds) mock_load_data.return_value = (mock_train, mock_test) _, _ = cifar100_dataset.get_centralized_datasets() mock_load_data.assert_called_once() # Assert the validation ClientData isn't used, and the train and test # are amalgamated into datasets single datasets over all clients. 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())
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, crop_size: Optional[int] = 24, max_batches: Optional[int] = None, cache_dir: Optional[str] = '~'): """Trains a ResNet-18 on CIFAR-100 using a given optimizer. 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. crop_size: The crop size used for CIFAR-100 preprocessing. 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. """ crop_shape = (crop_size, crop_size, NUM_CHANNELS) cifar_train, cifar_test = cifar100_dataset.get_centralized_datasets( train_batch_size=batch_size, crop_shape=crop_shape, cache_dir=cache_dir) if max_batches and max_batches >= 1: cifar_train = cifar_train.take(max_batches) cifar_test = cifar_test.take(max_batches) model = resnet_models.create_resnet18(input_shape=crop_shape, num_classes=NUM_CLASSES) model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(), optimizer=optimizer, metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]) centralized_training_loop.run(keras_model=model, train_dataset=cifar_train, validation_dataset=cifar_test, 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 test_centralized_cifar_structure(self): crop_shape = (24, 24, 3) cifar_train, cifar_test = cifar100_dataset.get_centralized_datasets( train_batch_size=20, test_batch_size=100, crop_shape=crop_shape) train_batch = next(iter(cifar_train)) train_batch_shape = tuple(train_batch[0].shape) self.assertEqual(train_batch_shape, (20, 24, 24, 3)) test_batch = next(iter(cifar_test)) test_batch_shape = tuple(test_batch[0].shape) self.assertEqual(test_batch_shape, (100, 24, 24, 3))
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 configure_training( task_spec: training_specs.TaskSpec, crop_size: int = 24, distort_train_images: bool = True) -> training_specs.RunnerSpec: """Configures training for the CIFAR-100 classification task. 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. crop_size: An optional integer representing the resulting size of input images after preprocessing. distort_train_images: A boolean indicating whether to distort training images during preprocessing via random crops, as opposed to simply resizing the image. Returns: A `RunnerSpec` containing attributes used for running the newly created federated task. """ crop_shape = (crop_size, crop_size, 3) cifar_train, _ = tff.simulation.datasets.cifar100.load_data() _, cifar_test = cifar100_dataset.get_centralized_datasets( train_batch_size=task_spec.client_batch_size, crop_shape=crop_shape) train_preprocess_fn = cifar100_dataset.create_preprocess_fn( num_epochs=task_spec.client_epochs_per_round, batch_size=task_spec.client_batch_size, crop_shape=crop_shape, distort_image=distort_train_images) input_spec = train_preprocess_fn.type_signature.result.element 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()) iterative_process = task_spec.iterative_process_builder(tff_model_fn) @tff.tf_computation(tf.string) def build_train_dataset_from_client_id(client_id): client_dataset = cifar_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) client_ids_fn = training_utils.build_sample_fn( cifar_train.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_eval_fn = training_utils.build_centralized_evaluate_fn( eval_dataset=cifar_test, model_builder=model_builder, loss_builder=loss_builder, metrics_builder=metrics_builder) def test_fn(state): return centralized_eval_fn(iterative_process.get_model_weights(state)) def validation_fn(state, round_num): del round_num return test_fn(state) return training_specs.RunnerSpec(iterative_process=training_process, client_datasets_fn=client_sampling_fn, validation_fn=validation_fn, test_fn=test_fn)
def test_raises_length_2_crop(self): with self.assertRaises(ValueError): cifar100_dataset.get_federated_datasets(crop_shape=(32, 32)) with self.assertRaises(ValueError): cifar100_dataset.get_centralized_datasets(crop_shape=(32, 32))
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, 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. 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. 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, _ = cifar100_dataset.get_federated_datasets( train_client_epochs_per_round=client_epochs_per_round, train_client_batch_size=client_batch_size, crop_shape=crop_shape) _, cifar_test = cifar100_dataset.get_centralized_datasets( train_batch_size=client_batch_size, crop_shape=crop_shape) if max_eval_batches and max_eval_batches >= 1: cifar_test = cifar_test.take(max_eval_batches) 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) client_datasets_fn = training_utils.build_client_datasets_fn( dataset=cifar_train, clients_per_round=clients_per_round, random_seed=client_datasets_random_seed) evaluate_fn = training_utils.build_centralized_evaluate_fn( eval_dataset=cifar_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)