experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add distribution strategy tests.

Browse Source 
PiperOrigin-RevId: 308090534
Change-Id: I70c08479f4ae324f3d7a8835bf89b0ac10105811
This commit is contained in:
A. Unique TensorFlower 2020-04-23 11:09:55 -07:00 committed by TensorFlower Gardener
parent 958d2c10e6
commit 8052c3b3f6
11 changed files with 690 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
tensorflow/python/keras/distribute/BUILD
View File

@ -440,3 +440,15 @@ py_library(
"//tensorflow/python/keras/optimizer_v2", "//tensorflow/python/keras/optimizer_v2",
], ],
) )
py_library(
name = "tpu_strategy_test_utils",
srcs = ["tpu_strategy_test_utils.py"],
deps = [
"//tensorflow/python:platform",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/eager:remote",
"//tensorflow/python/tpu:tpu_strategy_util",
],
)

47
tensorflow/python/keras/distribute/tpu_strategy_test_utils.py Normal file
View File

@ -0,0 +1,47 @@
# Copyright 2020 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.
# ==============================================================================
"""Utility functions for tests using TPUStrategy."""
from __future__ import absolute_import
from __future__ import division
from __future__ import google_type_annotations
from __future__ import print_function
from tensorflow.python.distribute import tpu_strategy
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
from tensorflow.python.eager import remote
from tensorflow.python.platform import flags
from tensorflow.python.tpu import tpu_strategy_util
FLAGS = flags.FLAGS
flags.DEFINE_string("tpu", "", "Name of TPU to connect to.")
flags.DEFINE_string("project", None, "Name of GCP project with TPU.")
flags.DEFINE_string("zone", None, "Name of GCP zone with TPU.")
def get_tpu_cluster_resolver():
resolver = tpu_cluster_resolver.TPUClusterResolver(
tpu=FLAGS.tpu,
zone=FLAGS.zone,
project=FLAGS.project,
)
return resolver
def get_tpu_strategy():
resolver = get_tpu_cluster_resolver()
remote.connect_to_cluster(resolver)
tpu_strategy_util.initialize_tpu_system(resolver)
return tpu_strategy.TPUStrategy(resolver)

143
tensorflow/python/keras/layers/preprocessing/BUILD
View File

@ -3,6 +3,7 @@
load("//tensorflow:tensorflow.bzl", "tf_py_test") load("//tensorflow:tensorflow.bzl", "tf_py_test")
load("//tensorflow:tensorflow.bzl", "cuda_py_test") load("//tensorflow:tensorflow.bzl", "cuda_py_test")
load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_test")
package( package(
default_visibility = [ default_visibility = [
@ -231,19 +232,6 @@ py_library(
], ],
) )
tf_py_test(
name = "discretization_test",
size = "small",
srcs = ["discretization_test.py"],
python_version = "PY3",
deps = [
":discretization",
":preprocessing_test_utils",
"//tensorflow/python:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_test( cuda_py_test(
name = "categorical_crossing_test", name = "categorical_crossing_test",
size = "medium", size = "medium",
@ -261,6 +249,63 @@ cuda_py_test(
], ],
) )
tf_py_test(
name = "categorical_encoding_test",
size = "medium",
srcs = ["categorical_encoding_test.py"],
python_version = "PY3",
deps = [
":categorical_encoding",
":preprocessing_test_utils",
"//tensorflow/python:client_testlib",
"//tensorflow/python/keras",
"//tensorflow/python/keras/utils:generic_utils",
"//tensorflow/python/ops/ragged:ragged_string_ops",
"@absl_py//absl/testing:parameterized",
],
)
tpu_py_test(
name = "categorical_encoding_tpu_test",
srcs = ["categorical_encoding_tpu_test.py"],
disable_experimental = True,
python_version = "PY3",
tags = ["no_oss"],
deps = [
":categorical_encoding",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
],
)
tf_py_test(
name = "discretization_test",
size = "small",
srcs = ["discretization_test.py"],
python_version = "PY3",
deps = [
":discretization",
":preprocessing_test_utils",
"//tensorflow/python:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
tpu_py_test(
name = "discretization_tpu_test",
srcs = ["discretization_tpu_test.py"],
disable_experimental = True,
python_version = "PY3",
tags = ["no_oss"],
deps = [
":discretization",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
],
)
cuda_py_test( cuda_py_test(
name = "hashing_test", name = "hashing_test",
size = "medium", size = "medium",
@ -275,6 +320,20 @@ cuda_py_test(
], ],
) )
tpu_py_test(
name = "hashing_tpu_test",
srcs = ["hashing_tpu_test.py"],
disable_experimental = True,
python_version = "PY3",
tags = ["no_oss"],
deps = [
":hashing",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
],
)
tf_py_test( tf_py_test(
name = "index_lookup_test", name = "index_lookup_test",
size = "medium", size = "medium",
@ -291,6 +350,20 @@ tf_py_test(
], ],
) )
tpu_py_test(
name = "index_lookup_tpu_test",
srcs = ["index_lookup_tpu_test.py"],
disable_experimental = True,
python_version = "PY3",
tags = ["no_oss"],
deps = [
":index_lookup",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
],
)
cuda_py_test( cuda_py_test(
name = "image_preprocessing_test", name = "image_preprocessing_test",
size = "medium", size = "medium",
@ -318,6 +391,20 @@ tf_py_test(
], ],
) )
tpu_py_test(
name = "normalization_tpu_test",
srcs = ["normalization_tpu_test.py"],
disable_experimental = True,
python_version = "PY3",
tags = ["no_oss"],
deps = [
":normalization",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
],
)
tf_py_test( tf_py_test(
name = "text_vectorization_test", name = "text_vectorization_test",
size = "medium", size = "medium",
@ -334,6 +421,20 @@ tf_py_test(
], ],
) )
tpu_py_test(
name = "text_vectorization_tpu_test",
srcs = ["text_vectorization_tpu_test.py"],
disable_experimental = True,
python_version = "PY3",
tags = ["no_oss"],
deps = [
":text_vectorization",
"//tensorflow/python/eager:test",
"//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
],
)
tf_py_test( tf_py_test(
name = "reduction_test", name = "reduction_test",
srcs = ["reduction_test.py"], srcs = ["reduction_test.py"],
@ -345,22 +446,6 @@ tf_py_test(
], ],
) )
tf_py_test(
name = "categorical_encoding_test",
size = "medium",
srcs = ["categorical_encoding_test.py"],
python_version = "PY3",
deps = [
":categorical_encoding",
":preprocessing_test_utils",
"//tensorflow/python:client_testlib",
"//tensorflow/python/keras",
"//tensorflow/python/keras/utils:generic_utils",
"//tensorflow/python/ops/ragged:ragged_string_ops",
"@absl_py//absl/testing:parameterized",
],
)
tf_py_test( tf_py_test(
name = "preprocessing_stage_test", name = "preprocessing_stage_test",
srcs = ["preprocessing_stage_test.py"], srcs = ["preprocessing_stage_test.py"],

60
tensorflow/python/keras/layers/preprocessing/categorical_encoding_tpu_test.py Normal file
View File

@ -0,0 +1,60 @@
# Copyright 2019 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.
# ==============================================================================
"""Tests for keras.layers.preprocessing.normalization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensorflow.python import keras
from tensorflow.python.framework import dtypes
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.distribute import tpu_strategy_test_utils
from tensorflow.python.keras.layers.preprocessing import categorical_encoding
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes(
always_skip_v1=True, always_skip_eager=True)
class CategoricalEncodingDistributionTest(
keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self):
input_array = np.array([[1, 2, 3, 1], [0, 3, 1, 0]])
# pyformat: disable
expected_output = [[0, 1, 1, 1, 0, 0],
[1, 1, 0, 1, 0, 0]]
# pyformat: enable
max_tokens = 6
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope():
input_data = keras.Input(shape=(4,), dtype=dtypes.int32)
layer = categorical_encoding.CategoricalEncoding(
max_tokens=max_tokens, output_mode=categorical_encoding.BINARY)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_array)
self.assertAllEqual(expected_output, output_dataset)
if __name__ == "__main__":
test.main()

57
tensorflow/python/keras/layers/preprocessing/discretization_tpu_test.py Normal file
View File

@ -0,0 +1,57 @@
# Copyright 2019 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.
# ==============================================================================
"""Tests for keras.layers.preprocessing.normalization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensorflow.python import keras
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.distribute import tpu_strategy_test_utils
from tensorflow.python.keras.layers.preprocessing import discretization
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes(
always_skip_v1=True, always_skip_eager=True)
class DiscretizationDistributionTest(
keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self):
input_array = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])
expected_output = [[0, 2, 3, 1], [1, 3, 2, 1]]
expected_output_shape = [None, None]
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope():
input_data = keras.Input(shape=(None,))
layer = discretization.Discretization(
bins=[0., 1., 2.], output_mode=discretization.INTEGER)
bucket_data = layer(input_data)
self.assertAllEqual(expected_output_shape, bucket_data.shape.as_list())
model = keras.Model(inputs=input_data, outputs=bucket_data)
output_dataset = model.predict(input_array)
self.assertAllEqual(expected_output, output_dataset)
if __name__ == "__main__":
test.main()

58
tensorflow/python/keras/layers/preprocessing/hashing_tpu_test.py Normal file
View File

@ -0,0 +1,58 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for keras.layers.preprocessing.normalization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import config
from tensorflow.python.framework import dtypes
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.distribute import tpu_strategy_test_utils
from tensorflow.python.keras.layers.preprocessing import hashing
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes(
always_skip_v1=True, always_skip_eager=True)
class HashingDistributionTest(keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self):
input_data = np.asarray([["omar"], ["stringer"], ["marlo"], ["wire"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_data).batch(
2, drop_remainder=True)
expected_output = [[0], [0], [1], [0]]
config.set_soft_device_placement(True)
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = hashing.Hashing(num_bins=2)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
if __name__ == "__main__":
test.main()

24
tensorflow/python/keras/layers/preprocessing/index_lookup_test.py
View File

@ -30,6 +30,7 @@ from tensorflow.python import keras
from tensorflow.python import tf2 from tensorflow.python import tf2
from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import one_device_strategy
from tensorflow.python.eager import context from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes from tensorflow.python.framework import dtypes
from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import sparse_tensor
@ -403,6 +404,29 @@ class CategoricalEncodingAdaptTest(
layer.adapt(batched_ds) layer.adapt(batched_ds)
@keras_parameterized.run_all_keras_modes
class IndexLookupDistributionTest(
keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
def test_cpu_distribution(self):
vocab_data = ["earth", "wind", "and", "fire"]
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
strategy = one_device_strategy.OneDeviceStrategy("/cpu:0")
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = get_layer_class()()
layer.set_vocabulary(vocab_data)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_array)
self.assertAllEqual(expected_output, output_dataset)
@keras_parameterized.run_all_keras_modes @keras_parameterized.run_all_keras_modes
class IndexLookupOutputTest(keras_parameterized.TestCase, class IndexLookupOutputTest(keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest): preprocessing_test_utils.PreprocessingLayerTest):

66
tensorflow/python/keras/layers/preprocessing/index_lookup_tpu_test.py Normal file
View File

@ -0,0 +1,66 @@
# Copyright 2019 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.
# ==============================================================================
"""Tests for keras.layers.preprocessing.normalization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import config
from tensorflow.python.framework import dtypes
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.distribute import tpu_strategy_test_utils
from tensorflow.python.keras.layers.preprocessing import index_lookup
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes(
always_skip_v1=True, always_skip_eager=True)
class IndexLookupDistributionTest(
keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self):
vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = index_lookup.IndexLookup()
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
if __name__ == "__main__":
test.main()

128
tensorflow/python/keras/layers/preprocessing/normalization_tpu_test.py Normal file
View File

@ -0,0 +1,128 @@
# Copyright 2019 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.
# ==============================================================================
"""Tests for keras.layers.preprocessing.normalization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from absl.testing import parameterized
import numpy as np
from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.distribute import tpu_strategy_test_utils
from tensorflow.python.keras.layers.preprocessing import normalization
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test
def _get_layer_computation_test_cases():
test_cases = ({
"adapt_data": np.array([[1.], [2.], [3.], [4.], [5.]], dtype=np.float32),
"axis": -1,
"test_data": np.array([[1.], [2.], [3.]], np.float32),
"expected": np.array([[-1.414214], [-.707107], [0]], np.float32),
"testcase_name": "2d_single_element"
}, {
"adapt_data": np.array([[1.], [2.], [3.], [4.], [5.]], dtype=np.float32),
"axis": None,
"test_data": np.array([[1.], [2.], [3.]], np.float32),
"expected": np.array([[-1.414214], [-.707107], [0]], np.float32),
"testcase_name": "2d_single_element_none_axis"
}, {
"adapt_data": np.array([[1., 2., 3., 4., 5.]], dtype=np.float32),
"axis": None,
"test_data": np.array([[1.], [2.], [3.]], np.float32),
"expected": np.array([[-1.414214], [-.707107], [0]], np.float32),
"testcase_name": "2d_single_element_none_axis_flat_data"
}, {
"adapt_data":
np.array([[[1., 2., 3.], [2., 3., 4.]], [[3., 4., 5.], [4., 5., 6.]]],
np.float32),
"axis":
1,
"test_data":
np.array([[[1., 2., 3.], [2., 3., 4.]], [[3., 4., 5.], [4., 5., 6.]]],
np.float32),
"expected":
np.array([[[-1.549193, -0.774597, 0.], [-1.549193, -0.774597, 0.]],
[[0., 0.774597, 1.549193], [0., 0.774597, 1.549193]]],
np.float32),
"testcase_name":
"3d_internal_axis"
}, {
"adapt_data":
np.array(
[[[1., 0., 3.], [2., 3., 4.]], [[3., -1., 5.], [4., 5., 8.]]],
np.float32),
"axis": (1, 2),
"test_data":
np.array(
[[[3., 1., -1.], [2., 5., 4.]], [[3., 0., 5.], [2., 5., 8.]]],
np.float32),
"expected":
np.array(
[[[1., 3., -5.], [-1., 1., -1.]], [[1., 1., 1.], [-1., 1., 1.]]],
np.float32),
"testcase_name":
"3d_multiple_axis"
})
crossed_test_cases = []
# Cross above test cases with use_dataset in (True, False)
for use_dataset in (True, False):
for case in test_cases:
case = case.copy()
if use_dataset:
case["testcase_name"] = case["testcase_name"] + "_with_dataset"
case["use_dataset"] = use_dataset
crossed_test_cases.append(case)
return crossed_test_cases
@keras_parameterized.run_all_keras_modes(
always_skip_v1=True, always_skip_eager=True)
class NormalizationTest(keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
@parameterized.named_parameters(*_get_layer_computation_test_cases())
def test_layer_computation(self, adapt_data, axis, test_data, use_dataset,
expected):
input_shape = tuple([None for _ in range(test_data.ndim - 1)])
if use_dataset:
# Keras APIs expect batched datasets
adapt_data = dataset_ops.Dataset.from_tensor_slices(adapt_data).batch(
test_data.shape[0] // 2)
test_data = dataset_ops.Dataset.from_tensor_slices(test_data).batch(
test_data.shape[0] // 2)
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope():
input_data = keras.Input(shape=input_shape)
layer = normalization.Normalization(axis=axis)
layer.adapt(adapt_data)
output = layer(input_data)
model = keras.Model(input_data, output)
output_data = model.predict(test_data)
self.assertAllClose(expected, output_data)
if __name__ == "__main__":
test.main()

28
tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py
View File

@ -27,6 +27,7 @@ from tensorflow.python import keras
from tensorflow.python import tf2 from tensorflow.python import tf2
from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import one_device_strategy
from tensorflow.python.eager import context from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes from tensorflow.python.framework import dtypes
from tensorflow.python.keras import backend from tensorflow.python.keras import backend
@ -568,6 +569,33 @@ class TextVectorizationPreprocessingTest(
self.assertAllEqual(expected_output, output) self.assertAllEqual(expected_output, output)
@keras_parameterized.run_all_keras_modes
class TextVectorizationDistributionTest(
keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
def test_distribution_strategy_output(self):
vocab_data = ["earth", "wind", "and", "fire"]
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
strategy = one_device_strategy.OneDeviceStrategy("/cpu:0")
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = get_layer_class()(
max_tokens=None,
standardize=None,
split=None,
output_mode=text_vectorization.INT)
layer.set_vocabulary(vocab_data)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_array)
self.assertAllEqual(expected_output, output_dataset)
@keras_parameterized.run_all_keras_modes @keras_parameterized.run_all_keras_modes
class TextVectorizationOutputTest( class TextVectorizationOutputTest(
keras_parameterized.TestCase, keras_parameterized.TestCase,

96
tensorflow/python/keras/layers/preprocessing/text_vectorization_tpu_test.py Normal file
View File

@ -0,0 +1,96 @@
# Copyright 2019 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.
# ==============================================================================
"""Tests for keras.layers.preprocessing.normalization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import config
from tensorflow.python.framework import dtypes
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.distribute import tpu_strategy_test_utils
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.keras.layers.preprocessing import text_vectorization
from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes(
always_skip_v1=True, always_skip_eager=True)
class TextVectorizationTPUDistributionTest(
keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest):
def test_distribution_strategy_output(self):
vocab_data = ["earth", "wind", "and", "fire"]
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = text_vectorization.TextVectorization(
max_tokens=None,
standardize=None,
split=None,
output_mode=text_vectorization.INT)
layer.set_vocabulary(vocab_data)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def test_distribution_strategy_output_with_adapt(self):
vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = text_vectorization.TextVectorization(
max_tokens=None,
standardize=None,
split=None,
output_mode=text_vectorization.INT)
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
if __name__ == "__main__":
test.main()