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

Add distribution strategy tests for Keras preprocessing layers.

Browse Source 
PiperOrigin-RevId: 309027642
Change-Id: Id38df0a26c9e134034f9b8ceadc7f9b91c219158
This commit is contained in:
A. Unique TensorFlower 2020-04-29 08:29:59 -07:00 committed by TensorFlower Gardener
parent 01913de741
commit 4b848524a4
5 changed files with 69 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -266,17 +266,19 @@ tf_py_test(
], ],
) )
tpu_py_test( distribute_py_test(
name = "categorical_encoding_tpu_test", name = "categorical_encoding_distribution_test",
srcs = ["categorical_encoding_tpu_test.py"], srcs = ["categorical_encoding_distribution_test.py"],
disable_experimental = True, main = "categorical_encoding_distribution_test.py",
python_version = "PY3", python_version = "PY3",
tags = ["no_oss"], tags = [
"multi_and_single_gpu",
],
deps = [ deps = [
":categorical_encoding", ":categorical_encoding",
"//tensorflow/python/distribute:tpu_strategy", "//tensorflow/python/distribute:combinations",
"//tensorflow/python/distribute:strategy_combinations",
"//tensorflow/python/keras", "//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
], ],
) )
@ -293,17 +295,17 @@ tf_py_test(
], ],
) )
tpu_py_test( distribute_py_test(
name = "discretization_tpu_test", name = "discretization_distribution_test",
srcs = ["discretization_tpu_test.py"], srcs = ["discretization_distribution_test.py"],
disable_experimental = True, main = "discretization_distribution_test.py",
python_version = "PY3", python_version = "PY3",
tags = ["no_oss"], tags = ["multi_and_single_gpu"],
deps = [ deps = [
":discretization", ":discretization",
"//tensorflow/python/distribute:tpu_strategy", "//tensorflow/python/distribute:combinations",
"//tensorflow/python/distribute:strategy_combinations",
"//tensorflow/python/keras", "//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
], ],
) )
@ -322,16 +324,16 @@ cuda_py_test(
) )
tpu_py_test( tpu_py_test(
name = "hashing_tpu_test", name = "hashing_distribution_test",
srcs = ["hashing_tpu_test.py"], srcs = ["hashing_distribution_test.py"],
disable_experimental = True, main = "hashing_distribution_test.py",
python_version = "PY3", python_version = "PY3",
tags = ["no_oss"], tags = ["multi_and_single_gpu"],
deps = [ deps = [
":hashing", ":hashing",
"//tensorflow/python/distribute:tpu_strategy", "//tensorflow/python/distribute:combinations",
"//tensorflow/python/distribute:strategy_combinations",
"//tensorflow/python/keras", "//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
], ],
) )
@ -352,16 +354,16 @@ tf_py_test(
) )
tpu_py_test( tpu_py_test(
name = "index_lookup_tpu_test", name = "index_lookup_distribution_test",
srcs = ["index_lookup_tpu_test.py"], srcs = ["index_lookup_distribution_test.py"],
disable_experimental = True, main = "index_lookup_distribution_test.py",
python_version = "PY3", python_version = "PY3",
tags = ["no_oss"], tags = ["no_oss"],
deps = [ deps = [
":index_lookup", ":index_lookup",
"//tensorflow/python/distribute:tpu_strategy", "//tensorflow/python/distribute:combinations",
"//tensorflow/python/distribute:strategy_combinations",
"//tensorflow/python/keras", "//tensorflow/python/keras",
"//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
], ],
) )

15
tensorflow/python/keras/layers/preprocessing/categorical_encoding_tpu_test.py → tensorflow/python/keras/layers/preprocessing/categorical_encoding_distribution_test.py
View File

@ -21,21 +21,24 @@ from __future__ import print_function
import numpy as np import numpy as np
from tensorflow.python import keras from tensorflow.python import keras
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.framework import dtypes from tensorflow.python.framework import dtypes
from tensorflow.python.keras import keras_parameterized 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 categorical_encoding
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes( @combinations.generate(
always_skip_v1=True, always_skip_eager=True) combinations.combine(
distribution=strategy_combinations.all_strategies,
mode=["eager", "graph"]))
class CategoricalEncodingDistributionTest( class CategoricalEncodingDistributionTest(
keras_parameterized.TestCase, keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest): preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self): def test_distribution(self, distribution):
input_array = np.array([[1, 2, 3, 1], [0, 3, 1, 0]]) input_array = np.array([[1, 2, 3, 1], [0, 3, 1, 0]])
# pyformat: disable # pyformat: disable
@ -44,9 +47,7 @@ class CategoricalEncodingDistributionTest(
# pyformat: enable # pyformat: enable
max_tokens = 6 max_tokens = 6
strategy = tpu_strategy_test_utils.get_tpu_strategy() with distribution.scope():
with strategy.scope():
input_data = keras.Input(shape=(4,), dtype=dtypes.int32) input_data = keras.Input(shape=(4,), dtype=dtypes.int32)
layer = categorical_encoding.CategoricalEncoding( layer = categorical_encoding.CategoricalEncoding(
max_tokens=max_tokens, output_mode=categorical_encoding.BINARY) max_tokens=max_tokens, output_mode=categorical_encoding.BINARY)

14
tensorflow/python/keras/layers/preprocessing/discretization_tpu_test.py → tensorflow/python/keras/layers/preprocessing/discretization_distribution_test.py
View File

@ -21,27 +21,29 @@ from __future__ import print_function
import numpy as np import numpy as np
from tensorflow.python import keras from tensorflow.python import keras
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.keras import keras_parameterized 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 discretization
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes( @combinations.generate(
always_skip_v1=True, always_skip_eager=True) combinations.combine(
distribution=strategy_combinations.all_strategies,
mode=["eager", "graph"]))
class DiscretizationDistributionTest( class DiscretizationDistributionTest(
keras_parameterized.TestCase, keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest): preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self): def test_distribution(self, distribution):
input_array = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]]) 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 = [[0, 2, 3, 1], [1, 3, 2, 1]]
expected_output_shape = [None, None] expected_output_shape = [None, None]
strategy = tpu_strategy_test_utils.get_tpu_strategy() with distribution.scope():
with strategy.scope():
input_data = keras.Input(shape=(None,)) input_data = keras.Input(shape=(None,))
layer = discretization.Discretization( layer = discretization.Discretization(
bins=[0., 1., 2.], output_mode=discretization.INTEGER) bins=[0., 1., 2.], output_mode=discretization.INTEGER)

14
tensorflow/python/keras/layers/preprocessing/hashing_tpu_test.py → tensorflow/python/keras/layers/preprocessing/hashing_distribution_test.py
View File

@ -22,30 +22,32 @@ import numpy as np
from tensorflow.python import keras from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.framework import config from tensorflow.python.framework import config
from tensorflow.python.framework import dtypes from tensorflow.python.framework import dtypes
from tensorflow.python.keras import keras_parameterized 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 hashing
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes( @combinations.generate(
always_skip_v1=True, always_skip_eager=True) combinations.combine(
distribution=strategy_combinations.all_strategies,
mode=["eager", "graph"]))
class HashingDistributionTest(keras_parameterized.TestCase, class HashingDistributionTest(keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest): preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self): def test_distribution(self, distribution):
input_data = np.asarray([["omar"], ["stringer"], ["marlo"], ["wire"]]) input_data = np.asarray([["omar"], ["stringer"], ["marlo"], ["wire"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_data).batch( input_dataset = dataset_ops.Dataset.from_tensor_slices(input_data).batch(
2, drop_remainder=True) 2, drop_remainder=True)
expected_output = [[0], [0], [1], [0]] expected_output = [[0], [0], [1], [0]]
config.set_soft_device_placement(True) config.set_soft_device_placement(True)
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope(): with distribution.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string) input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = hashing.Hashing(num_bins=2) layer = hashing.Hashing(num_bins=2)
int_data = layer(input_data) int_data = layer(input_data)

25
tensorflow/python/keras/layers/preprocessing/index_lookup_tpu_test.py → tensorflow/python/keras/layers/preprocessing/index_lookup_distribution_test.py
View File

@ -22,22 +22,34 @@ import numpy as np
from tensorflow.python import keras from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.eager import context
from tensorflow.python.framework import config from tensorflow.python.framework import config
from tensorflow.python.framework import dtypes from tensorflow.python.framework import dtypes
from tensorflow.python.keras import keras_parameterized 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 index_lookup
from tensorflow.python.keras.layers.preprocessing import index_lookup_v1
from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
from tensorflow.python.platform import test from tensorflow.python.platform import test
@keras_parameterized.run_all_keras_modes( def get_layer_class():
always_skip_v1=True, always_skip_eager=True) if context.executing_eagerly():
return index_lookup.IndexLookup
else:
return index_lookup_v1.IndexLookup
@combinations.generate(
combinations.combine(
distribution=strategy_combinations.all_strategies,
mode=["eager", "graph"]))
class IndexLookupDistributionTest( class IndexLookupDistributionTest(
keras_parameterized.TestCase, keras_parameterized.TestCase,
preprocessing_test_utils.PreprocessingLayerTest): preprocessing_test_utils.PreprocessingLayerTest):
def test_tpu_distribution(self): def test_tpu_distribution(self, distribution):
vocab_data = [[ vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and", "earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire" "and", "fire"
@ -50,11 +62,10 @@ class IndexLookupDistributionTest(
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]] expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True) config.set_soft_device_placement(True)
strategy = tpu_strategy_test_utils.get_tpu_strategy()
with strategy.scope(): with distribution.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string) input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = index_lookup.IndexLookup() layer = get_layer_class()()
layer.adapt(vocab_dataset) layer.adapt(vocab_dataset)
int_data = layer(input_data) int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data) model = keras.Model(inputs=input_data, outputs=int_data)