def test_run_simple_model(self): vocab_size = 6 mask_model = shakespeare_models.create_recurrent_model( vocab_size, sequence_length=5) mask_model.compile(optimizer='sgd', loss='sparse_categorical_crossentropy', metrics=[keras_metrics.MaskedCategoricalAccuracy()]) no_mask_model = shakespeare_models.create_recurrent_model( vocab_size, sequence_length=5, mask_zero=False) no_mask_model.compile( optimizer='sgd', loss='sparse_categorical_crossentropy', metrics=[keras_metrics.MaskedCategoricalAccuracy()]) constant_test_weights = tf.nest.map_structure(tf.ones_like, mask_model.weights) mask_model.set_weights(constant_test_weights) no_mask_model.set_weights(constant_test_weights) # `tf.data.Dataset.from_tensor_slices` aggresively coalesces the input into # a single tensor, but we want a tuple of two tensors per example, so we # apply a transformation to split. def split_to_tuple(t): return (t[0, :], t[1, :]) data = tf.data.Dataset.from_tensor_slices([ ([0, 1, 2, 3, 4], [1, 2, 3, 4, 0]), ([2, 3, 4, 0, 1], [3, 4, 0, 1, 2]), ]).map(split_to_tuple).batch(2) mask_metrics = mask_model.evaluate(data) no_mask_metrics = no_mask_model.evaluate(data) self.assertNotAllClose(mask_metrics, no_mask_metrics, atol=1e-3)
def test_model_initialization_uses_random_seed(self): model_1_with_seed_0 = shakespeare_models.create_recurrent_model( vocab_size=6, sequence_length=5, seed=0) model_2_with_seed_0 = shakespeare_models.create_recurrent_model( vocab_size=6, sequence_length=5, seed=0) model_1_with_seed_1 = shakespeare_models.create_recurrent_model( vocab_size=6, sequence_length=5, seed=1) model_2_with_seed_1 = shakespeare_models.create_recurrent_model( vocab_size=6, sequence_length=5, seed=1) self.assertAllClose(model_1_with_seed_0.weights, model_2_with_seed_0.weights) self.assertAllClose(model_1_with_seed_1.weights, model_2_with_seed_1.weights) self.assertNotAllClose(model_1_with_seed_0.weights, model_1_with_seed_1.weights)
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 create_shakespeare_model(sequence_length): """Constructs a `tf.keras.Model` to train.""" return shakespeare_models.create_recurrent_model( vocab_size=VOCAB_SIZE, sequence_length=sequence_length)