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