Module keras.utils.kpl_test_utils
Test related utilities for KPL + tf.distribute.
Expand source code
# 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 tensorflow.compat.v2 as tf
import random
import tempfile
import keras
from keras.layers.preprocessing import string_lookup
class DistributeKplTestUtils(tf.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=tf.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=tf.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 = tf.data.Dataset.from_generator(
feature_and_label_gen,
output_signature={
"features": tf.TensorSpec([3], tf.string),
"label": tf.TensorSpec([1], tf.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=tf.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 = tf.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."""
@tf.function
def serve_fn(raw_features):
raw_features = tf.expand_dims(raw_features, axis=0)
transformed_features = model.feature_mapper(raw_features)
outputs = model(transformed_features)
outputs = tf.squeeze(outputs, axis=0)
outputs = tf.cast(tf.greater(outputs, 0.5), tf.int64)
decoded_outputs = model.label_inverse_lookup_layer(outputs)
return tf.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(
tf.TensorSpec(
shape=(3), dtype=tf.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(
tf.constant(["avenger", "ironman", "avenger"]))["output_0"]
self.assertIn(prediction0.numpy().decode("UTF-8"), ("yes", "no"))
prediction1 = loaded_serving_fn(
tf.constant(["ironman", "ironman", "unkonwn"]))["output_0"]
self.assertIn(prediction1.numpy().decode("UTF-8"), ("yes", "no"))
Classes
class DistributeKplTestUtils (methodName='runTest')
-
Utils for test of tf.distribute + KPL.
Create an instance of the class that will use the named test method when executed. Raises a ValueError if the instance does not have a method with the specified name.
Expand source code
class DistributeKplTestUtils(tf.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=tf.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=tf.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 = tf.data.Dataset.from_generator( feature_and_label_gen, output_signature={ "features": tf.TensorSpec([3], tf.string), "label": tf.TensorSpec([1], tf.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=tf.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 = tf.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.""" @tf.function def serve_fn(raw_features): raw_features = tf.expand_dims(raw_features, axis=0) transformed_features = model.feature_mapper(raw_features) outputs = model(transformed_features) outputs = tf.squeeze(outputs, axis=0) outputs = tf.cast(tf.greater(outputs, 0.5), tf.int64) decoded_outputs = model.label_inverse_lookup_layer(outputs) return tf.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( tf.TensorSpec( shape=(3), dtype=tf.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( tf.constant(["avenger", "ironman", "avenger"]))["output_0"] self.assertIn(prediction0.numpy().decode("UTF-8"), ("yes", "no")) prediction1 = loaded_serving_fn( tf.constant(["ironman", "ironman", "unkonwn"]))["output_0"] self.assertIn(prediction1.numpy().decode("UTF-8"), ("yes", "no"))
Ancestors
- tensorflow.python.framework.test_util.TensorFlowTestCase
- absl.testing.absltest.TestCase
- absl.third_party.unittest3_backport.case.TestCase
- unittest.case.TestCase
Class variables
var FEATURE_VOCAB
var LABEL_VOCAB
Methods
def create_serving_signature(self, model, feature_mapper, label_inverse_lookup_layer)
-
Create serving signature for the given model.
Expand source code
def create_serving_signature(self, model, feature_mapper, label_inverse_lookup_layer): """Create serving signature for the given model.""" @tf.function def serve_fn(raw_features): raw_features = tf.expand_dims(raw_features, axis=0) transformed_features = model.feature_mapper(raw_features) outputs = model(transformed_features) outputs = tf.squeeze(outputs, axis=0) outputs = tf.cast(tf.greater(outputs, 0.5), tf.int64) decoded_outputs = model.label_inverse_lookup_layer(outputs) return tf.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( tf.TensorSpec( shape=(3), dtype=tf.string, name="example"))
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.
Expand source code
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 = tf.data.Dataset.from_generator( feature_and_label_gen, output_signature={ "features": tf.TensorSpec([3], tf.string), "label": tf.TensorSpec([1], tf.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_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.
Expand source code
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=tf.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=tf.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 define_model(self)
-
A simple model for test of tf.distribute + KPL.
Expand source code
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=tf.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 = tf.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.
Expand source code
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 test_save_load_serving_model(self, model, feature_mapper, label_inverse_lookup_layer)
-
Test save/load/serving model.
Expand source code
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( tf.constant(["avenger", "ironman", "avenger"]))["output_0"] self.assertIn(prediction0.numpy().decode("UTF-8"), ("yes", "no")) prediction1 = loaded_serving_fn( tf.constant(["ironman", "ironman", "unkonwn"]))["output_0"] self.assertIn(prediction1.numpy().decode("UTF-8"), ("yes", "no"))