| # Copyright 2021 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Test related utilities for KPL + tf.distribute.""" |
| |
| import random |
| import tempfile |
| |
| from tensorflow.python import keras |
| from tensorflow.python.data.ops import dataset_ops |
| from tensorflow.python.eager import def_function |
| from tensorflow.python.framework import constant_op |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import tensor_spec |
| from tensorflow.python.keras.layers.preprocessing import string_lookup |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import math_ops |
| from tensorflow.python.platform import test |
| |
| |
| class DistributeKplTestUtils(test.TestCase): |
| """Utils for test of tf.distribute + KPL.""" |
| FEATURE_VOCAB = [ |
| "avenger", "ironman", "batman", "hulk", "spiderman", "kingkong", |
| "wonder_woman" |
| ] |
| LABEL_VOCAB = ["yes", "no"] |
| |
| def define_kpls_for_training(self, use_adapt): |
| """Function that defines KPL used for unit tests of tf.distribute. |
| |
| Args: |
| use_adapt: if adapt will be called. False means there will be precomputed |
| statistics. |
| |
| Returns: |
| feature_mapper: a simple keras model with one keras StringLookup layer |
| which maps feature to index. |
| label_mapper: similar to feature_mapper, but maps label to index. |
| |
| """ |
| if use_adapt: |
| feature_lookup_layer = ( |
| string_lookup.StringLookup( |
| num_oov_indices=1)) |
| feature_lookup_layer.adapt(self.FEATURE_VOCAB) |
| label_lookup_layer = ( |
| string_lookup.StringLookup( |
| num_oov_indices=0, mask_token=None)) |
| label_lookup_layer.adapt(self.LABEL_VOCAB) |
| else: |
| feature_lookup_layer = ( |
| string_lookup.StringLookup( |
| vocabulary=self.FEATURE_VOCAB, num_oov_indices=1)) |
| label_lookup_layer = ( |
| string_lookup.StringLookup( |
| vocabulary=self.LABEL_VOCAB, num_oov_indices=0, mask_token=None)) |
| |
| raw_feature_input = keras.layers.Input( |
| shape=(3,), dtype=dtypes.string, name="feature", ragged=True) |
| feature_id_input = feature_lookup_layer(raw_feature_input) |
| feature_mapper = keras.Model({"features": raw_feature_input}, |
| feature_id_input) |
| |
| raw_label_input = keras.layers.Input( |
| shape=(1,), dtype=dtypes.string, name="label") |
| label_id_input = label_lookup_layer(raw_label_input) |
| label_mapper = keras.Model({"label": raw_label_input}, label_id_input) |
| |
| return feature_mapper, label_mapper |
| |
| def dataset_fn(self, feature_mapper, label_mapper): |
| """Function that generates dataset for test of tf.distribute + KPL. |
| |
| Args: |
| feature_mapper: a simple keras model with one keras StringLookup layer |
| which maps feature to index. |
| label_mapper: similar to feature_mapper, but maps label to index. |
| |
| Returns: |
| Generated dataset for test of tf.distribute + KPL. |
| |
| """ |
| |
| def feature_and_label_gen(): |
| # Generator of dataset. |
| while True: |
| features = random.sample(self.FEATURE_VOCAB, 3) |
| label = ["yes"] if self.FEATURE_VOCAB[0] in features else ["no"] |
| yield {"features": features, "label": label} |
| |
| raw_dataset = dataset_ops.Dataset.from_generator( |
| feature_and_label_gen, |
| output_signature={ |
| "features": tensor_spec.TensorSpec([3], dtypes.string), |
| "label": tensor_spec.TensorSpec([1], dtypes.string) |
| }).shuffle(100).batch(32) |
| |
| train_dataset = raw_dataset.map(lambda x: ( # pylint: disable=g-long-lambda |
| { |
| "features": feature_mapper(x["features"]) |
| }, label_mapper(x["label"]))) |
| return train_dataset |
| |
| def define_model(self): |
| """A simple model for test of tf.distribute + KPL.""" |
| # Create the model. The input needs to be compatible with KPLs. |
| model_input = keras.layers.Input( |
| shape=(3,), dtype=dtypes.int64, name="model_input") |
| |
| # input_dim includes a mask token and an oov token. |
| emb_output = keras.layers.Embedding( |
| input_dim=len(self.FEATURE_VOCAB) + 2, output_dim=20)( |
| model_input) |
| emb_output = math_ops.reduce_mean(emb_output, axis=1) |
| dense_output = keras.layers.Dense( |
| units=1, activation="sigmoid")( |
| emb_output) |
| model = keras.Model({"features": model_input}, dense_output) |
| return model |
| |
| def define_reverse_lookup_layer(self): |
| """Create string reverse lookup layer for serving.""" |
| |
| label_inverse_lookup_layer = string_lookup.StringLookup( |
| num_oov_indices=0, |
| mask_token=None, |
| vocabulary=self.LABEL_VOCAB, |
| invert=True) |
| return label_inverse_lookup_layer |
| |
| def create_serving_signature(self, model, feature_mapper, |
| label_inverse_lookup_layer): |
| """Create serving signature for the given model.""" |
| |
| @def_function.function |
| def serve_fn(raw_features): |
| raw_features = array_ops.expand_dims(raw_features, axis=0) |
| transformed_features = model.feature_mapper(raw_features) |
| outputs = model(transformed_features) |
| outputs = array_ops.squeeze(outputs, axis=0) |
| outputs = math_ops.cast(math_ops.greater(outputs, 0.5), dtypes.int64) |
| decoded_outputs = model.label_inverse_lookup_layer(outputs) |
| return array_ops.squeeze(decoded_outputs, axis=0) |
| |
| model.feature_mapper = feature_mapper |
| model.label_inverse_lookup_layer = label_inverse_lookup_layer |
| # serving does NOT have batch dimension |
| return serve_fn.get_concrete_function( |
| tensor_spec.TensorSpec( |
| shape=(3), dtype=dtypes.string, name="example")) |
| |
| def test_save_load_serving_model(self, model, feature_mapper, |
| label_inverse_lookup_layer): |
| """Test save/load/serving model.""" |
| |
| serving_fn = self.create_serving_signature(model, feature_mapper, |
| label_inverse_lookup_layer) |
| |
| saved_model_dir = tempfile.mkdtemp(dir=self.get_temp_dir()) |
| model.save(saved_model_dir, save_format="tf", |
| signatures={"serving_default": serving_fn}) |
| |
| # Test the saved_model. |
| loaded_serving_fn = keras.saving.save.load_model( |
| saved_model_dir).signatures["serving_default"] |
| |
| # check the result w/ and w/o avenger. |
| prediction0 = loaded_serving_fn( |
| constant_op.constant(["avenger", "ironman", "avenger"]))["output_0"] |
| self.assertIn(prediction0.numpy().decode("UTF-8"), ("yes", "no")) |
| |
| prediction1 = loaded_serving_fn( |
| constant_op.constant(["ironman", "ironman", "unkonwn"]))["output_0"] |
| self.assertIn(prediction1.numpy().decode("UTF-8"), ("yes", "no")) |
