def metrics_builder():
"""Returns a `list` of `tf.keras.metric.Metric` objects."""
pad_token, _, _, _ = dataset.get_special_tokens()
return [
keras_metrics.NumBatchesCounter(),
keras_metrics.NumExamplesCounter(),
keras_metrics.NumTokensCounter(masked_tokens=[pad_token]),
keras_metrics.MaskedCategoricalAccuracy(masked_tokens=[pad_token]),
]
def test_to_ids(self):
pad, _, bos, eos = dataset.get_special_tokens()
to_tokens = dataset._build_tokenize_fn(split_length=5)
tokens = self.evaluate(to_tokens({'snippets': tf.constant('abc')}))
self.assertAllEqual(tokens, [bos, 64, 42, 21, eos])
to_tokens = dataset._build_tokenize_fn(split_length=12)
tokens = self.evaluate(
to_tokens({'snippets': tf.constant('star wars')}))
self.assertAllEqual(tokens,
[bos, 25, 5, 64, 46, 14, 26, 64, 46, 25, eos, pad])
def test_convert_snippets_to_character_sequence_examples(self):
pad, _, bos, eos = dataset.get_special_tokens()
ds = dataset.convert_snippets_to_character_sequence_examples(
tf.data.Dataset.from_tensor_slices({
'snippets': ['a snippet', 'different snippet'],
}),
batch_size=2,
epochs=2,
shuffle_buffer_size=1,
sequence_length=10)
expected_outputs = [
# First batch.
([[bos, 64, 14, 25, 45, 66, 4, 4, 65, 5],
[bos, 1, 66, 43, 43, 65, 46, 65, 45,
5]], [[64, 14, 25, 45, 66, 4, 4, 65, 5, eos],
[1, 66, 43, 43, 65, 46, 65, 45, 5, 14]]),
# Second batch.
([
[25, 45, 66, 4, 4, 65, 5, eos, pad, pad],
[bos, 64, 14, 25, 45, 66, 4, 4, 65, 5],
], [
[45, 66, 4, 4, 65, 5, eos, pad, pad, pad],
[64, 14, 25, 45, 66, 4, 4, 65, 5, eos],
]),
# Third batch.
([[bos, 1, 66, 43, 43, 65, 46, 65, 45, 5],
[25, 45, 66, 4, 4, 65, 5, eos, pad,
pad]], [[1, 66, 43, 43, 65, 46, 65, 45, 5, 14],
[45, 66, 4, 4, 65, 5, eos, pad, pad, pad]]),
]
for batch_num, actual in enumerate(ds):
self.assertGreater(
len(expected_outputs),
0,
msg='Actual output contains more than expected.\nActual: {!s}'.
format(actual))
expected = expected_outputs.pop(0)
self.assertAllEqual(
actual,
expected,
msg='Batch {:d} not equal. Actual: {!s}\nExpected: {!s}'.
format(batch_num, actual, expected))
self.assertLen(
expected_outputs,
0,
msg='Not all expected output seen.\nLeft over expectations: {!s}'.
format(expected_outputs))
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
experiment_output_dir = FLAGS.root_output_dir
tensorboard_dir = os.path.join(experiment_output_dir, 'logdir',
FLAGS.experiment_name)
results_dir = os.path.join(experiment_output_dir, 'results',
FLAGS.experiment_name)
for path in [experiment_output_dir, tensorboard_dir, results_dir]:
try:
tf.io.gfile.makedirs(path)
except tf.errors.OpError:
pass # Directory already exists.
hparam_dict = collections.OrderedDict([(name, FLAGS[name].value)
for name in hparam_flags])
hparam_dict['results_file'] = results_dir
hparams_file = os.path.join(results_dir, 'hparams.csv')
logging.info('Saving hyper parameters to: [%s]', hparams_file)
utils_impl.atomic_write_to_csv(pd.Series(hparam_dict), hparams_file)
train_client_data, test_client_data = (
tff.simulation.datasets.shakespeare.load_data())
def preprocess(ds):
return dataset.convert_snippets_to_character_sequence_examples(
ds, FLAGS.batch_size, epochs=1).cache()
train_dataset = train_client_data.create_tf_dataset_from_all_clients()
if FLAGS.shuffle_train_data:
train_dataset = train_dataset.shuffle(buffer_size=10000)
train_dataset = preprocess(train_dataset)
eval_dataset = preprocess(
test_client_data.create_tf_dataset_from_all_clients())
optimizer = optimizer_utils.create_optimizer_fn_from_flags('centralized')()
pad_token, _, _, _ = dataset.get_special_tokens()
# Vocabulary with one OOV ID and zero for the mask.
vocab_size = len(dataset.CHAR_VOCAB) + 2
model = models.create_recurrent_model(vocab_size=vocab_size,
batch_size=FLAGS.batch_size)
model.compile(
optimizer=optimizer,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[
keras_metrics.MaskedCategoricalAccuracy(masked_tokens=[pad_token])
])
logging.info('Training model:')
logging.info(model.summary())
csv_logger_callback = keras_callbacks.AtomicCSVLogger(results_dir)
tensorboard_callback = tf.keras.callbacks.TensorBoard(
log_dir=tensorboard_dir)
# Reduce the learning rate every 20 epochs.
def decay_lr(epoch, lr):
if (epoch + 1) % 20 == 0:
return lr * 0.1
else:
return lr
lr_callback = tf.keras.callbacks.LearningRateScheduler(decay_lr, verbose=1)
history = model.fit(
train_dataset,
validation_data=eval_dataset,
epochs=FLAGS.num_epochs,
callbacks=[lr_callback, tensorboard_callback, csv_logger_callback])
logging.info('Final metrics:')
for name in ['loss', 'accuracy']:
metric = history.history['val_{}'.format(name)][-1]
logging.info('\t%s: %.4f', name, metric)
