def test_create_client_optimizer_from_flags(self, optimizer_name, optimizer_cls): with flag_sandbox( {'{}_optimizer'.format(TEST_CLIENT_FLAG_PREFIX): optimizer_name}): # Construct a default optimizer. default_optimizer = utils_impl.create_optimizer_from_flags( TEST_CLIENT_FLAG_PREFIX) self.assertIsInstance(default_optimizer, optimizer_cls) # Override the default flag value. overridden_learning_rate = 5.0 custom_optimizer = utils_impl.create_optimizer_from_flags( TEST_CLIENT_FLAG_PREFIX, overrides={'learning_rate': overridden_learning_rate}) self.assertIsInstance(custom_optimizer, optimizer_cls) self.assertEqual(custom_optimizer.get_config()['learning_rate'], overridden_learning_rate) # Override learning rate flag. commandline_set_learning_rate = 100.0 with flag_sandbox({ '{}_learning_rate'.format(TEST_CLIENT_FLAG_PREFIX): commandline_set_learning_rate }): custom_optimizer = utils_impl.create_optimizer_from_flags( TEST_CLIENT_FLAG_PREFIX) self.assertIsInstance(custom_optimizer, optimizer_cls) self.assertEqual( custom_optimizer.get_config()['learning_rate'], commandline_set_learning_rate)
def create_compiled_keras_model(only_digits=True): """Create compiled keras model based on the original FedAvg CNN.""" data_format = 'channels_last' input_shape = [28, 28, 1] model = tf.keras.models.Sequential([ tf.keras.layers.Reshape(input_shape=(28 * 28, ), target_shape=input_shape), tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape, data_format=data_format), tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation='relu', data_format=data_format), tf.keras.layers.MaxPool2D(pool_size=(2, 2), data_format=data_format), tf.keras.layers.Dropout(0.25), tf.keras.layers.Flatten(), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(10 if only_digits else 62, activation=tf.nn.softmax), ]) model.compile(loss=tf.keras.losses.sparse_categorical_crossentropy, optimizer=utils_impl.create_optimizer_from_flags('client'), metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]) return model
def test_create_optimizer_from_flags_invalid_overrides(self): with flag_sandbox( {'{}_optimizer'.format(TEST_CLIENT_FLAG_PREFIX): 'sgd'}): with self.assertRaisesRegex(TypeError, 'type `collections.Mapping`'): _ = utils_impl.create_optimizer_from_flags( TEST_CLIENT_FLAG_PREFIX, overrides=[1, 2, 3])
def create_compiled_keras_model(): """Create compiled keras model based on the original FedAvg CNN.""" model = models.create_original_fedavg_cnn_model() model.compile(loss=tf.keras.losses.sparse_categorical_crossentropy, optimizer=utils_impl.create_optimizer_from_flags('client'), metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]) return model
def create_compiled_keras_model(): """Create compiled keras model.""" model = models.create_original_fedavg_cnn_model( only_digits=FLAGS.digit_only_emnist) model.compile(loss=tf.keras.losses.sparse_categorical_crossentropy, optimizer=utils_impl.create_optimizer_from_flags('client'), metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]) return model
def test_create_optimizer_from_flags_flags_set_not_for_optimizer(self): with flag_sandbox({'{}_optimizer'.format(TEST_CLIENT_FLAG_PREFIX): 'sgd'}): # Set an Adam flag that isn't used in SGD. # We need to use `_parse_args` because that is the only way FLAGS is # notified that a non-default value is being used. bad_adam_flag = '{}_adam_beta_1'.format(TEST_CLIENT_FLAG_PREFIX) FLAGS._parse_args( args=['--{}=0.5'.format(bad_adam_flag)], known_only=True) with self.assertRaisesRegex( ValueError, r'Commandline flags for .*\[sgd\].*\'test_client_adam_beta_1\'.*'): _ = utils_impl.create_optimizer_from_flags(TEST_CLIENT_FLAG_PREFIX) FLAGS[bad_adam_flag].unparse()
def model_fn(): """Defines the model.""" keras_model = model_builder() train_metrics = [ metrics.NumTokensCounter(name='num_tokens', masked_tokens=[pad]), metrics.NumTokensCounter(name='num_tokens_no_oov', masked_tokens=[pad, oov]), metrics.NumBatchesCounter(), metrics.NumExamplesCounter(), metrics.MaskedCategoricalAccuracy(name='accuracy', masked_tokens=[pad]), metrics.MaskedCategoricalAccuracy(name='accuracy_no_oov', masked_tokens=[pad, oov]), metrics.MaskedCategoricalAccuracy(name='accuracy_no_oov_no_eos', masked_tokens=[pad, oov, eos]), ] keras_model.compile( loss=tf.keras.losses.SparseCategoricalCrossentropy( from_logits=True), optimizer=utils_impl.create_optimizer_from_flags('client'), metrics=train_metrics) return tff.learning.from_compiled_keras_model(keras_model, sample_batch)
def run_experiment(): """Data preprocessing and experiment execution.""" emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data() example_tuple = collections.namedtuple('Example', ['x', 'y']) def element_fn(element): return example_tuple(x=tf.reshape(element['pixels'], [-1]), y=tf.reshape(element['label'], [1])) def preprocess_train_dataset(dataset): """Preprocess training dataset.""" return dataset.map(element_fn).apply( tf.data.experimental.shuffle_and_repeat( buffer_size=10000, count=FLAGS.client_epochs_per_round)).batch(FLAGS.batch_size) def preprocess_test_dataset(dataset): """Preprocess testing dataset.""" return dataset.map(element_fn).batch(100, drop_remainder=False) emnist_train = emnist_train.preprocess(preprocess_train_dataset) emnist_test = preprocess_test_dataset( emnist_test.create_tf_dataset_from_all_clients()) example_dataset = emnist_train.create_tf_dataset_for_client( emnist_train.client_ids[0]) sample_batch = tf.nest.map_structure(lambda x: x.numpy(), next(iter(example_dataset))) def model_fn(): keras_model = create_compiled_keras_model() return tff.learning.from_compiled_keras_model(keras_model, sample_batch) def client_datasets_fn(round_num): """Returns a list of client datasets.""" del round_num # Unused. sampled_clients = np.random.choice(emnist_train.client_ids, size=FLAGS.train_clients_per_round, replace=False) return [ emnist_train.create_tf_dataset_for_client(client) for client in sampled_clients ] tf.io.gfile.makedirs(FLAGS.root_output_dir) hparam_dict = collections.OrderedDict([(name, FLAGS[name].value) for name in hparam_flags]) the_metrics_hook = metrics_hook.MetricsHook.build( FLAGS.exp_name, FLAGS.root_output_dir, emnist_test, hparam_dict, create_compiled_keras_model()) optimizer_fn = lambda: utils_impl.create_optimizer_from_flags('server') if FLAGS.use_compression: # We create a `StatefulBroadcastFn` and `StatefulAggregateFn` 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 model_fn, and return instances of appropriate encoders. encoded_broadcast_fn = ( tff.learning.framework.build_encoded_broadcast_from_model( model_fn, _broadcast_encoder_fn)) encoded_mean_fn = tff.learning.framework.build_encoded_mean_from_model( model_fn, _mean_encoder_fn) else: encoded_broadcast_fn = None encoded_mean_fn = None training_loops.federated_averaging_training_loop( model_fn, optimizer_fn, client_datasets_fn, total_rounds=FLAGS.total_rounds, rounds_per_eval=FLAGS.rounds_per_eval, metrics_hook=the_metrics_hook, stateful_model_broadcast_fn=encoded_broadcast_fn, stateful_delta_aggregate_fn=encoded_mean_fn)
def _run_experiment(): """Data preprocessing and experiment execution.""" emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data( only_digits=FLAGS.digit_only_emnist) def preprocess_train_dataset(dataset): """Preprocess training dataset.""" return (dataset.map(reshape_emnist_element).shuffle( buffer_size=10000).repeat(FLAGS.client_epochs_per_round).batch( FLAGS.batch_size)) def preprocess_test_dataset(dataset): """Preprocess testing dataset.""" return dataset.map(reshape_emnist_element).batch(100, drop_remainder=False) emnist_train = emnist_train.preprocess(preprocess_train_dataset) emnist_test = preprocess_test_dataset( emnist_test.create_tf_dataset_from_all_clients()) example_dataset = emnist_train.create_tf_dataset_for_client( emnist_train.client_ids[0]) input_spec = example_dataset.element_spec 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()) def client_datasets_fn(round_num): """Returns a list of client datasets.""" del round_num # Unused. sampled_clients = random.sample(population=emnist_train.client_ids, k=FLAGS.train_clients_per_round) return [ emnist_train.create_tf_dataset_for_client(client) for client in sampled_clients ] def evaluate_fn(state): compiled_keras_model = compiled_eval_keras_model() tff.learning.assign_weights_to_keras_model(compiled_keras_model, state.model) eval_metrics = compiled_keras_model.evaluate(emnist_test, verbose=0) return { 'loss': eval_metrics[0], 'sparse_categorical_accuracy': eval_metrics[1], } tf.io.gfile.makedirs(FLAGS.root_output_dir) hparam_dict = collections.OrderedDict([(name, FLAGS[name].value) for name in hparam_flags]) hparam_dict = utils_impl.remove_unused_flags('client', hparam_dict) metrics_hook = _MetricsHook(FLAGS.exp_name, FLAGS.root_output_dir, hparam_dict) client_optimizer_fn = lambda: utils_impl.create_optimizer_from_flags( 'client') if FLAGS.server_optimizer == 'sgd': server_optimizer_fn = functools.partial( tf.keras.optimizers.SGD, learning_rate=FLAGS.server_learning_rate, momentum=FLAGS.server_momentum) elif FLAGS.server_optimizer == 'flars': server_optimizer_fn = functools.partial( flars_optimizer.FLARSOptimizer, learning_rate=FLAGS.server_learning_rate, momentum=FLAGS.server_momentum, max_ratio=FLAGS.max_ratio) else: raise ValueError('Optimizer %s is not supported.' % FLAGS.server_optimizer) _federated_averaging_training_loop(model_fn=tff_model_fn, client_optimizer_fn=client_optimizer_fn, server_optimizer_fn=server_optimizer_fn, client_datasets_fn=client_datasets_fn, evaluate_fn=evaluate_fn, total_rounds=FLAGS.total_rounds, rounds_per_eval=FLAGS.rounds_per_eval, metrics_hook=metrics_hook)
def run_experiment(): """Runs the training experiment.""" try: tf.io.gfile.makedirs( os.path.join(FLAGS.root_output_dir, FLAGS.exp_name)) except tf.errors.OpError: pass train_set, validation_set, test_set = ( dataset.construct_word_level_datasets( vocab_size=FLAGS.vocab_size, client_batch_size=FLAGS.batch_size, client_epochs_per_round=1, max_seq_len=FLAGS.sequence_length, max_elements_per_user=FLAGS.max_elements_per_user, centralized_train=True, shuffle_buffer_size=None, num_validation_examples=FLAGS.num_validation_examples, num_test_examples=FLAGS.num_test_examples)) recurrent_model = tf.keras.layers.LSTM if FLAGS.lstm else tf.keras.layers.GRU def _layer_fn(): return recurrent_model(FLAGS.latent_size, return_sequences=True) pad, oov, _, eos = dataset.get_special_tokens(FLAGS.vocab_size) model = models.create_recurrent_model( FLAGS.vocab_size, FLAGS.embedding_size, FLAGS.num_layers, _layer_fn, 'stackoverflow-recurrent', shared_embedding=FLAGS.shared_embedding) logging.info('Training model: %s', model.summary()) optimizer = utils_impl.create_optimizer_from_flags('centralized') model.compile( loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=optimizer, metrics=[ metrics.MaskedCategoricalAccuracy([pad], 'accuracy_with_oov'), metrics.MaskedCategoricalAccuracy([pad, oov], 'accuracy_no_oov'), metrics.MaskedCategoricalAccuracy([pad, oov, eos], 'accuracy_no_oov_no_eos') ]) train_results_path = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'train_results') test_results_path = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'test_results') train_csv_logger = AtomicCSVLogger(train_results_path) test_csv_logger = AtomicCSVLogger(test_results_path) log_dir = os.path.join(FLAGS.root_output_dir, 'logdir', FLAGS.exp_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) results_file = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'results.csv.bz2') # Write the hyperparameters to a CSV: hparam_dict = collections.OrderedDict([(name, FLAGS[name].value) for name in hparam_flags]) hparam_dict['results_file'] = results_file hparams_file = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'hparams.csv') utils_impl.atomic_write_to_csv(pd.Series(hparam_dict), hparams_file) model.fit(train_set, epochs=FLAGS.epochs, verbose=1, steps_per_epoch=FLAGS.steps_per_epoch, validation_data=validation_set, callbacks=[train_csv_logger, train_tensorboard_callback]) score = model.evaluate( test_set, verbose=1, callbacks=[test_csv_logger, test_tensorboard_callback]) logging.info('Final test loss: %.4f', score[0]) logging.info('Final test accuracy: %.4f', score[1])
def test_create_optimizer_from_flags_invalid_optimizer(self): FLAGS['{}_optimizer'.format(TEST_CLIENT_FLAG_PREFIX)].value = 'foo' with self.assertRaisesRegex(ValueError, 'not a valid optimizer'): _ = utils_impl.create_optimizer_from_flags(TEST_CLIENT_FLAG_PREFIX)
def run_experiment(): """Runs the training experiment.""" training_set, validation_set, test_set = ( dataset.construct_word_level_datasets( vocab_size=FLAGS.vocab_size, batch_size=FLAGS.batch_size, client_epochs_per_round=1, max_seq_len=FLAGS.sequence_length, max_training_elements_per_user=-1, num_validation_examples=FLAGS.num_validation_examples, num_test_examples=FLAGS.num_test_examples)) centralized_train = training_set.create_tf_dataset_from_all_clients() def _lstm_fn(): return tf.keras.layers.LSTM(FLAGS.latent_size, return_sequences=True) model = models.create_recurrent_model( FLAGS.vocab_size, FLAGS.embedding_size, FLAGS.num_layers, _lstm_fn, 'stackoverflow-lstm', shared_embedding=FLAGS.shared_embedding) logging.info('Training model: %s', model.summary()) optimizer = utils_impl.create_optimizer_from_flags('centralized') model.compile(loss=tf.keras.losses.sparse_categorical_crossentropy, optimizer=optimizer, weighted_metrics=['acc']) train_results_path = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'train_results') test_results_path = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'test_results') train_csv_logger = AtomicCSVLogger(train_results_path) test_csv_logger = AtomicCSVLogger(test_results_path) log_dir = os.path.join(FLAGS.root_output_dir, 'logdir', FLAGS.exp_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) results_file = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'results.csv.bz2') # Write the hyperparameters to a CSV: hparam_dict = collections.OrderedDict([(name, FLAGS[name].value) for name in hparam_flags]) hparam_dict['results_file'] = results_file hparams_file = os.path.join(FLAGS.root_output_dir, FLAGS.exp_name, 'hparams.csv') utils_impl.atomic_write_to_csv(pd.Series(hparam_dict), hparams_file) oov, bos, eos, pad = dataset.get_special_tokens(FLAGS.vocab_size) class_weight = {x: 1.0 for x in range(FLAGS.vocab_size)} class_weight[oov] = 0.0 # No credit for predicting OOV. class_weight[bos] = 0.0 # Shouldn't matter since this is never a target. class_weight[eos] = 1.0 # Model should learn to predict end of sentence. class_weight[pad] = 0.0 # No credit for predicting pad. model.fit(centralized_train, epochs=FLAGS.epochs, verbose=1, class_weight=class_weight, validation_data=validation_set, callbacks=[train_csv_logger, train_tensorboard_callback]) score = model.evaluate( test_set, verbose=1, callbacks=[test_csv_logger, test_tensorboard_callback]) logging.info('Final test loss: %.4f', score[0]) logging.info('Final test accuracy: %.4f', score[1])
def server_optimizer_fn(): return utils_impl.create_optimizer_from_flags('server')