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')
