class DistributionStrategyLstmModelCorrectnessTest( _DistributionStrategyRnnModelCorrectnessTest): def _get_layer_class(self): if tf2.enabled(): if not context.executing_eagerly(): self.skipTest( "LSTM v2 and legacy graph mode don't work together.") return rnn_v2.LSTM else: return rnn_v1.LSTM @ds_combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model() + keras_correctness_test_base.multi_worker_mirrored_eager()) def test_lstm_model_correctness(self, distribution, use_numpy, use_validation_data): self.run_correctness_test(distribution, use_numpy, use_validation_data) @ds_combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model() + keras_correctness_test_base.multi_worker_mirrored_eager()) @testing_utils.enable_v2_dtype_behavior def test_lstm_model_correctness_mixed_precision(self, distribution, use_numpy, use_validation_data): if isinstance(distribution, (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1)): policy_name = 'mixed_bfloat16' else: policy_name = 'mixed_float16' with policy.policy_scope(policy_name): self.run_correctness_test(distribution, use_numpy, use_validation_data)
class DistributionStrategyEmbeddingModelCorrectnessTest( keras_correctness_test_base .TestDistributionStrategyEmbeddingModelCorrectnessBase): def get_model(self, max_words=10, initial_weights=None, distribution=None, run_distributed=None, input_shapes=None): del input_shapes with keras_correctness_test_base.MaybeDistributionScope(distribution): word_ids = keras.layers.Input( shape=(max_words,), dtype=np.int32, name='words') word_embed = keras.layers.Embedding(input_dim=20, output_dim=10)(word_ids) if self.use_distributed_dense: word_embed = keras.layers.TimeDistributed(keras.layers.Dense(4))( word_embed) avg = keras.layers.GlobalAveragePooling1D()(word_embed) preds = keras.layers.Dense(2, activation='softmax')(avg) model = keras.Model(inputs=[word_ids], outputs=[preds]) if initial_weights: model.set_weights(initial_weights) model.compile( # TODO(b/130808953): Switch back the V1 optimizer once global_step is # mirrored. optimizer=gradient_descent_keras.SGD(learning_rate=0.1), loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy'], run_distributed=run_distributed) return model @combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model()) def test_embedding_model_correctness(self, distribution, use_numpy, use_validation_data, run_distributed): self.use_distributed_dense = False self.run_correctness_test(distribution, use_numpy, use_validation_data, run_distributed) @combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model()) def test_embedding_time_distributed_model_correctness(self, distribution, use_numpy, use_validation_data, run_distributed): self.use_distributed_dense = True self.run_correctness_test(distribution, use_numpy, use_validation_data, run_distributed)
class DistributionStrategyLstmModelCorrectnessTest( keras_correctness_test_base. TestDistributionStrategyEmbeddingModelCorrectnessBase): def get_model(self, max_words=10, initial_weights=None, distribution=None): with keras_correctness_test_base.MaybeDistributionScope(distribution): word_ids = keras.layers.Input(shape=(max_words, ), dtype=np.int32, name='words') word_embed = keras.layers.Embedding(input_dim=20, output_dim=10)(word_ids) lstm_embed = keras.layers.LSTM(units=4, return_sequences=False)(word_embed) preds = keras.layers.Dense(2, activation='softmax')(lstm_embed) model = keras.Model(inputs=[word_ids], outputs=[preds]) if initial_weights: model.set_weights(initial_weights) model.compile(optimizer=gradient_descent.GradientDescentOptimizer( learning_rate=0.1), loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy']) return model @combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model()) def test_lstm_model_correctness(self, distribution, use_numpy, use_validation_data): self.run_correctness_test(distribution, use_numpy, use_validation_data)
class DistributionStrategyLstmModelCorrectnessTest( keras_correctness_test_base .TestDistributionStrategyEmbeddingModelCorrectnessBase): def get_model(self, max_words=10, initial_weights=None, distribution=None, experimental_run_tf_function=None, input_shapes=None): del input_shapes if tf2.enabled(): if not context.executing_eagerly(): self.skipTest("LSTM v2 and legacy graph mode don't work together.") lstm = rnn_v2.LSTM else: lstm = rnn_v1.LSTM with keras_correctness_test_base.MaybeDistributionScope(distribution): word_ids = keras.layers.Input( shape=(max_words,), dtype=np.int32, name='words') word_embed = keras.layers.Embedding(input_dim=20, output_dim=10)(word_ids) lstm_embed = lstm(units=4, return_sequences=False)( word_embed) preds = keras.layers.Dense(2, activation='softmax')(lstm_embed) model = keras.Model(inputs=[word_ids], outputs=[preds]) if initial_weights: model.set_weights(initial_weights) optimizer_fn = gradient_descent_keras.SGD model.compile( optimizer=optimizer_fn(learning_rate=0.1), loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy'], experimental_run_tf_function=experimental_run_tf_function) return model @combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model()) def test_lstm_model_correctness(self, distribution, use_numpy, use_validation_data, experimental_run_tf_function): self.run_correctness_test(distribution, use_numpy, use_validation_data, experimental_run_tf_function)
class DistributionStrategyGruModelCorrectnessTest( _DistributionStrategyRnnModelCorrectnessTest): def _get_layer_class(self): if tf2.enabled(): if not context.executing_eagerly(): self.skipTest( "GRU v2 and legacy graph mode don't work together.") return rnn_v2.GRU else: return rnn_v1.GRU @ds_combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model()) def test_gru_model_correctness(self, distribution, use_numpy, use_validation_data): self.run_correctness_test(distribution, use_numpy, use_validation_data)
class DistributionStrategyLstmModelCorrectnessTest( _DistributionStrategyRnnModelCorrectnessTest): def _get_layer_class(self): if tf2.enabled(): if not context.executing_eagerly(): self.skipTest("LSTM v2 and legacy graph mode don't work together.") return rnn_v2.LSTM else: return rnn_v1.LSTM @combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model()) def test_lstm_model_correctness(self, distribution, use_numpy, use_validation_data, experimental_run_tf_function): self.run_correctness_test(distribution, use_numpy, use_validation_data, experimental_run_tf_function)
class DistributionStrategySiameseEmbeddingModelCorrectnessTest( keras_correctness_test_base. TestDistributionStrategyEmbeddingModelCorrectnessBase): def get_model(self, max_words=10, initial_weights=None, distribution=None, input_shapes=None): del input_shapes with keras_correctness_test_base.MaybeDistributionScope(distribution): word_ids_a = keras.layers.Input(shape=(max_words, ), dtype=np.int32, name='words_a') word_ids_b = keras.layers.Input(shape=(max_words, ), dtype=np.int32, name='words_b') def submodel(embedding, word_ids): word_embed = embedding(word_ids) rep = keras.layers.GlobalAveragePooling1D()(word_embed) return keras.Model(inputs=[word_ids], outputs=[rep]) word_embed = keras.layers.Embedding( input_dim=20, output_dim=10, input_length=max_words, embeddings_initializer=keras.initializers.RandomUniform(0, 1)) a_rep = submodel(word_embed, word_ids_a).outputs[0] b_rep = submodel(word_embed, word_ids_b).outputs[0] sim = keras.layers.Dot(axes=1, normalize=True)([a_rep, b_rep]) model = keras.Model(inputs=[word_ids_a, word_ids_b], outputs=[sim]) if initial_weights: model.set_weights(initial_weights) # TODO(b/130808953): Switch back to the V1 optimizer after global_step # is made mirrored. model.compile( optimizer=gradient_descent_keras.SGD(learning_rate=0.1), loss='mse', metrics=['mse']) return model def get_data(self, count=(keras_correctness_test_base._GLOBAL_BATCH_SIZE * keras_correctness_test_base._EVAL_STEPS), min_words=5, max_words=10, max_word_id=19, num_classes=2): features_a, labels_a, _ = (super( DistributionStrategySiameseEmbeddingModelCorrectnessTest, self).get_data(count, min_words, max_words, max_word_id, num_classes)) features_b, labels_b, _ = (super( DistributionStrategySiameseEmbeddingModelCorrectnessTest, self).get_data(count, min_words, max_words, max_word_id, num_classes)) y_train = np.zeros((count, 1), dtype=np.float32) y_train[labels_a == labels_b] = 1.0 y_train[labels_a != labels_b] = -1.0 # TODO(b/123360757): Add tests for using list as inputs for multi-input # models. x_train = { 'words_a': features_a, 'words_b': features_b, } x_predict = x_train return x_train, y_train, x_predict @ds_combinations.generate( keras_correctness_test_base.test_combinations_for_embedding_model()) def test_siamese_embedding_model_correctness(self, distribution, use_numpy, use_validation_data): self.run_correctness_test(distribution, use_numpy, use_validation_data)