• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Test related utilities for KPL + tf.distribute."""
16from __future__ import absolute_import
17from __future__ import division
18from __future__ import print_function
19
20import random
21import tempfile
22
23from tensorflow.python import keras
24from tensorflow.python.data.ops import dataset_ops
25from tensorflow.python.eager import def_function
26from tensorflow.python.framework import constant_op
27from tensorflow.python.framework import dtypes
28from tensorflow.python.framework import tensor_spec
29from tensorflow.python.keras.layers.preprocessing import string_lookup
30from tensorflow.python.ops import array_ops
31from tensorflow.python.ops import math_ops
32from tensorflow.python.platform import test
33from tensorflow.python.saved_model import save as tf_save
34
35
36class DistributeKplTestUtils(test.TestCase):
37  """Utils for test of tf.distribute + KPL."""
38  FEATURE_VOCAB = [
39      "avenger", "ironman", "batman", "hulk", "spiderman", "kingkong",
40      "wonder_woman"
41  ]
42  LABEL_VOCAB = ["yes", "no"]
43
44  def define_kpls_for_training(self, use_adapt):
45    """Function that defines KPL used for unit tests of tf.distribute.
46
47    Args:
48      use_adapt: if adapt will be called. False means there will be precomputed
49        statistics.
50
51    Returns:
52      feature_mapper: a simple keras model with one keras StringLookup layer
53      which maps feature to index.
54      label_mapper: similar to feature_mapper, but maps label to index.
55
56    """
57    if use_adapt:
58      feature_lookup_layer = (
59          string_lookup.StringLookup(
60              num_oov_indices=1))
61      feature_lookup_layer.adapt(self.FEATURE_VOCAB)
62      label_lookup_layer = (
63          string_lookup.StringLookup(
64              num_oov_indices=0, mask_token=None))
65      label_lookup_layer.adapt(self.LABEL_VOCAB)
66    else:
67      feature_lookup_layer = (
68          string_lookup.StringLookup(
69              vocabulary=self.FEATURE_VOCAB, num_oov_indices=1))
70      label_lookup_layer = (
71          string_lookup.StringLookup(
72              vocabulary=self.LABEL_VOCAB, num_oov_indices=0, mask_token=None))
73
74    raw_feature_input = keras.layers.Input(
75        shape=(3,), dtype=dtypes.string, name="feature", ragged=True)
76    feature_id_input = feature_lookup_layer(raw_feature_input)
77    feature_mapper = keras.Model({"features": raw_feature_input},
78                                 feature_id_input)
79
80    raw_label_input = keras.layers.Input(
81        shape=(1,), dtype=dtypes.string, name="label")
82    label_id_input = label_lookup_layer(raw_label_input)
83    label_mapper = keras.Model({"label": raw_label_input}, label_id_input)
84
85    return feature_mapper, label_mapper
86
87  def dataset_fn(self, feature_mapper, label_mapper):
88    """Function that generates dataset for test of tf.distribute + KPL.
89
90    Args:
91      feature_mapper: a simple keras model with one keras StringLookup layer
92        which maps feature to index.
93      label_mapper: similar to feature_mapper, but maps label to index.
94
95    Returns:
96      Generated dataset for test of tf.distribute + KPL.
97
98    """
99
100    def feature_and_label_gen():
101      # Generator of dataset.
102      while True:
103        features = random.sample(self.FEATURE_VOCAB, 3)
104        label = ["yes"] if self.FEATURE_VOCAB[0] in features else ["no"]
105        yield {"features": features, "label": label}
106
107    raw_dataset = dataset_ops.Dataset.from_generator(
108        feature_and_label_gen,
109        output_signature={
110            "features": tensor_spec.TensorSpec([3], dtypes.string),
111            "label": tensor_spec.TensorSpec([1], dtypes.string)
112        }).shuffle(100).batch(32)
113
114    train_dataset = raw_dataset.map(lambda x: (  # pylint: disable=g-long-lambda
115        {
116            "features": feature_mapper(x["features"])
117        }, label_mapper(x["label"])))
118    return train_dataset
119
120  def define_model(self):
121    """A simple model for test of tf.distribute + KPL."""
122    # Create the model. The input needs to be compatible with KPLs.
123    model_input = keras.layers.Input(
124        shape=(3,), dtype=dtypes.int64, name="model_input")
125
126    # input_dim includes a mask token and an oov token.
127    emb_output = keras.layers.Embedding(
128        input_dim=len(self.FEATURE_VOCAB) + 2, output_dim=20)(
129            model_input)
130    emb_output = math_ops.reduce_mean(emb_output, axis=1)
131    dense_output = keras.layers.Dense(
132        units=1, activation="sigmoid")(
133            emb_output)
134    model = keras.Model({"features": model_input}, dense_output)
135    return model
136
137  def define_reverse_lookup_layer(self):
138    """Create string reverse lookup layer for serving."""
139
140    label_inverse_lookup_layer = string_lookup.StringLookup(
141        num_oov_indices=1,
142        mask_token=None,
143        vocabulary=self.LABEL_VOCAB,
144        invert=True)
145    return label_inverse_lookup_layer
146
147  def create_serving_signature(self, model, feature_mapper,
148                               label_inverse_lookup_layer):
149    """Create serving signature for the given model."""
150
151    @def_function.function
152    def serve_fn(raw_features):
153      raw_features = array_ops.expand_dims(raw_features, axis=0)
154      transformed_features = model.feature_mapper(raw_features)
155      outputs = model(transformed_features)
156      outputs = array_ops.squeeze(outputs, axis=0)
157      outputs = math_ops.cast(math_ops.greater(outputs, 0.5), dtypes.int64)
158      decoded_outputs = model.label_inverse_lookup_layer(outputs)
159      return array_ops.squeeze(decoded_outputs, axis=0)
160
161    model.feature_mapper = feature_mapper
162    model.label_inverse_lookup_layer = label_inverse_lookup_layer
163    # serving does NOT have batch dimension
164    return serve_fn.get_concrete_function(
165        tensor_spec.TensorSpec(
166            shape=(3), dtype=dtypes.string, name="example"))
167
168  def test_save_load_serving_model(self, model, feature_mapper,
169                                   label_inverse_lookup_layer):
170    """Test save/load/serving model."""
171
172    serving_fn = self.create_serving_signature(model, feature_mapper,
173                                               label_inverse_lookup_layer)
174
175    saved_model_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
176    tf_save.save(
177        model, saved_model_dir, signatures={"serving_default": serving_fn})
178
179    # Test the saved_model.
180    loaded_serving_fn = keras.saving.save.load_model(
181        saved_model_dir).signatures["serving_default"]
182
183    # check the result w/ and w/o avenger.
184    prediction0 = loaded_serving_fn(
185        constant_op.constant(["avenger", "ironman", "avenger"]))["output_0"]
186    self.assertIn(prediction0.numpy().decode("UTF-8"), ("yes", "no"))
187
188    prediction1 = loaded_serving_fn(
189        constant_op.constant(["ironman", "ironman", "unkonwn"]))["output_0"]
190    self.assertIn(prediction1.numpy().decode("UTF-8"), ("yes", "no"))
191