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