def eval_metrics_builder():
pad_token, _, _, _ = shakespeare_dataset.get_special_tokens()
return [
tf.keras.metrics.SparseCategoricalCrossentropy(),
keras_metrics.MaskedCategoricalAccuracy(masked_tokens=[pad_token]),
]
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,
sequence_length: Optional[int] = 80,
max_batches: Optional[int] = None):
"""Trains a two-layer RNN on Shakespeare next-character-prediction.
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.
sequence_length: The sequence length used for Shakespeare 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.
"""
shakespeare_train, shakespeare_test = shakespeare_dataset.get_centralized_datasets(
train_batch_size=batch_size, sequence_length=sequence_length)
if max_batches and max_batches >= 1:
shakespeare_train = shakespeare_train.take(max_batches)
shakespeare_test = shakespeare_test.take(max_batches)
pad_token, _, _, _ = shakespeare_dataset.get_special_tokens()
model = shakespeare_models.create_recurrent_model(
vocab_size=VOCAB_SIZE, sequence_length=sequence_length)
model.compile(
optimizer=optimizer,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[
keras_metrics.MaskedCategoricalAccuracy(masked_tokens=[pad_token])
])
centralized_training_loop.run(keras_model=model,
train_dataset=shakespeare_train,
validation_dataset=shakespeare_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_to_ids(self):
pad, _, bos, eos = shakespeare_dataset.get_special_tokens()
to_tokens = shakespeare_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 = shakespeare_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 metrics_builder():
"""Returns a `list` of `tf.keras.metric.Metric` objects."""
pad_token, _, _, _ = shakespeare_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_preprocess_fn(self):
pad, _, bos, eos = shakespeare_dataset.get_special_tokens()
initial_ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
snippets=['a snippet', 'different snippet']))
preprocess_fn = shakespeare_dataset.create_preprocess_fn(
num_epochs=2,
batch_size=2,
shuffle_buffer_size=1,
sequence_length=10)
ds = preprocess_fn(initial_ds)
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 test_last_id_not_oov(self):
_, oov, bos, eos = shakespeare_dataset.get_special_tokens()
to_tokens = shakespeare_dataset._build_tokenize_fn(split_length=5)
tokens = to_tokens({'snippets': tf.constant('a\r~')})
self.assertAllEqual(tokens, [bos, 64, 86, oov, eos])
