Cleanup keras distribution tests.

PiperOrigin-RevId: 345066746
Change-Id: Ibd5adc02e9b1ae7c9f4d640f6fa7ad8140de89ac

Author: qlzh727

tensorflow/python/keras/distribute/BUILD

@@ -69,6 +69,7 @@ py_library(
         ":model_combinations",
         ":multi_worker_testing_utils",
         ":saved_model_test_base",
+        ":test_example",
     ],
 )
 
@@ -202,7 +203,6 @@ cuda_py_test(
         "//tensorflow/python/distribute:multi_worker_test_base",
         "//tensorflow/python/distribute:multi_worker_util",
         "//tensorflow/python/distribute:strategy_combinations",
-        "//tensorflow/python/distribute:strategy_test_lib",
         "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/keras:testing_utils",
@@ -601,10 +601,10 @@ distribute_py_test(
         "notsan",
     ],
     deps = [
+        ":distribute_strategy_test_lib",
         ":keras_test_lib",
         ":optimizer_combinations",
         "//tensorflow/python/distribute:parameter_server_strategy",
-        "//tensorflow/python/keras/distribute:distribute_strategy_test_lib",
     ],
 )
 
@@ -654,6 +654,7 @@ distribute_py_test(
     ],
     deps = [
         ":optimizer_combinations",
+        ":test_example",
         "//tensorflow/python:control_flow_ops",
         "//tensorflow/python:control_flow_v2_toggles",
         "//tensorflow/python:framework_ops",
@@ -663,9 +664,7 @@ distribute_py_test(
         "//tensorflow/python/data/ops:dataset_ops",
         "//tensorflow/python/distribute:combinations",
         "//tensorflow/python/distribute:mirrored_strategy",
-        "//tensorflow/python/distribute:single_loss_example",
         "//tensorflow/python/distribute:strategy_combinations",
-        "//tensorflow/python/distribute:strategy_test_lib",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:test",
         "//tensorflow/python/ops/losses",
@@ -842,27 +841,6 @@ distribute_py_test(
     ],
 )
 
-distribute_py_test(
-    name = "step_fn_test",
-    srcs = ["step_fn_test.py"],
-    main = "step_fn_test.py",
-    tags = [
-        "multi_and_single_gpu",
-    ],
-    deps = [
-        ":optimizer_combinations",
-        "//tensorflow/python:framework_test_lib",
-        "//tensorflow/python:variables",
-        "//tensorflow/python/distribute:combinations",
-        "//tensorflow/python/distribute:single_loss_example",
-        "//tensorflow/python/distribute:strategy_combinations",
-        "//tensorflow/python/eager:context",
-        "//tensorflow/python/eager:test",
-        "//third_party/py/numpy",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)
-
 tf_py_test(
     name = "parameter_server_training_test",
     srcs = ["parameter_server_training_test.py"],
@@ -962,3 +940,16 @@ py_library(
         "//tensorflow/python/distribute:strategy_combinations",
     ],
 )
+
+py_library(
+    name = "test_example",
+    srcs = ["test_example.py"],
+    deps = [
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:constant_op",
+        "//tensorflow/python:framework_ops",
+        "//tensorflow/python:layers",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python/data/ops:dataset_ops",
+    ],
+)

tensorflow/python/keras/distribute/minimize_loss_test.py

@@ -25,20 +25,19 @@
 from tensorflow.python.distribute import combinations as ds_combinations
 from tensorflow.python.distribute import reduce_util
 from tensorflow.python.distribute import strategy_combinations
-from tensorflow.python.distribute import strategy_test_lib
-from tensorflow.python.distribute.single_loss_example import batchnorm_example
-from tensorflow.python.distribute.single_loss_example import minimize_loss_example
 from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_combinations as combinations
 from tensorflow.python.keras.distribute import optimizer_combinations
+from tensorflow.python.keras.distribute.test_example import batchnorm_example
+from tensorflow.python.keras.distribute.test_example import minimize_loss_example
+from tensorflow.python.keras.layers import core
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import control_flow_v2_toggles
 from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import variable_scope
 from tensorflow.python.ops import variables as variables_lib
 from tensorflow.python.ops.losses import losses_impl
@@ -212,7 +211,7 @@ def run_step():
       def get_expected_variables(num_parameter_devices):
         name = optimizer._name
 
-        if strategy_test_lib.is_optimizer_v2_instance(optimizer):
+        if isinstance(optimizer, optimizer_v2.OptimizerV2):
           variables = VAR_MAP_V2[name]
         else:
           variables = VAR_MAP_V1[name]
@@ -353,7 +352,7 @@ def loss_fn():
 
         optimizer = optimizer_fn()  # GradientDescent with 0.2 learning rate
 
-        if strategy_test_lib.is_optimizer_v2_instance(optimizer):
+        if isinstance(optimizer, optimizer_v2.OptimizerV2):
           return optimizer.minimize(loss_fn, [w])
         else:
           if use_callable_loss:
@@ -430,25 +429,20 @@ def dataset_fn():
         return dataset.batch(batch_size=1, drop_remainder=True)
 
       optimizer = optimizer_fn()
-      kernel = strategy_test_lib.create_variable_like_keras_layer(
-          "kernel", (1, 1), dtypes.float32)
-      bias = strategy_test_lib.create_variable_like_keras_layer(
-          "bias", (1,), dtypes.float32)
-      # layer = core.Dense(1, use_bias=True)
+      layer = core.Dense(1, use_bias=True)
 
       key1 = "foo"
       value1 = "bar"
 
       def model_fn(output_context, x):
         """A very simple model written by the user."""
         def loss_fn():
-          y = array_ops.reshape(nn_ops.bias_add(
-              math_ops.matmul(x, kernel), bias), []) - constant_op.constant(1.)
+          y = array_ops.reshape(layer(x), []) - constant_op.constant(1.)
           return y * y
 
-        if strategy_test_lib.is_optimizer_v2_instance(optimizer):
+        if isinstance(optimizer, optimizer_v2.OptimizerV2):
           train_op = optimizer.minimize(
-              loss_fn, lambda: [kernel, bias])
+              loss_fn, lambda: layer.trainable_variables)
         else:
           train_op = optimizer.minimize(loss_fn)
         loss = loss_fn()
@@ -517,8 +511,8 @@ def run_step():
       for _ in range(5):
         _, loss = run_step()
         losses.append(loss)
-        weights.append(self.evaluate(kernel))
-        biases.append(self.evaluate(bias))
+        weights.append(self.evaluate(layer.kernel))
+        biases.append(self.evaluate(layer.bias))
 
       loss_is_not_increasing = all(y <= x for x, y in zip(losses, losses[1:]))
       self.assertTrue(loss_is_not_increasing)

tensorflow/python/keras/distribute/step_fn_test.py

@@ -0,0 +1,98 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""A simple network to use in tests and examples."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.layers import core
+from tensorflow.python.layers import normalization
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+
+
+def minimize_loss_example(optimizer, use_bias=False, use_callable_loss=True):
+  """Example of non-distribution-aware legacy code."""
+
+  def dataset_fn():
+    dataset = dataset_ops.Dataset.from_tensors([[1.]]).repeat()
+    # TODO(isaprykin): batch with drop_remainder causes shapes to be
+    # fully defined for TPU.  Remove this when XLA supports dynamic shapes.
+    return dataset.batch(1, drop_remainder=True)
+
+  layer = core.Dense(1, use_bias=use_bias)
+
+  def model_fn(x):
+    """A very simple model written by the user."""
+
+    def loss_fn():
+      y = array_ops.reshape(layer(x), []) - constant_op.constant(1.)
+      return y * y
+
+    if isinstance(optimizer, optimizer_v2.OptimizerV2):
+      return optimizer.minimize(loss_fn, lambda: layer.trainable_variables)
+    elif use_callable_loss:
+      return optimizer.minimize(loss_fn)
+    else:
+      return optimizer.minimize(loss_fn())
+
+  return model_fn, dataset_fn, layer
+
+
+def batchnorm_example(optimizer_fn,
+                      batch_per_epoch=1,
+                      momentum=0.9,
+                      renorm=False,
+                      update_ops_in_replica_mode=False):
+  """Example of non-distribution-aware legacy code with batch normalization."""
+
+  def dataset_fn():
+    # input shape is [16, 8], input values are increasing in both dimensions.
+    return dataset_ops.Dataset.from_tensor_slices(
+        [[[float(x * 8 + y + z * 100)
+           for y in range(8)]
+          for x in range(16)]
+         for z in range(batch_per_epoch)]).repeat()
+
+  optimizer = optimizer_fn()
+  batchnorm = normalization.BatchNormalization(
+      renorm=renorm, momentum=momentum, fused=False)
+  layer = core.Dense(1, use_bias=False)
+
+  def model_fn(x):
+    """A model that uses batchnorm."""
+
+    def loss_fn():
+      y = batchnorm(x, training=True)
+      with ops.control_dependencies(
+          ops.get_collection(ops.GraphKeys.UPDATE_OPS)
+          if update_ops_in_replica_mode else []):
+        loss = math_ops.reduce_mean(
+            math_ops.reduce_sum(layer(y)) - constant_op.constant(1.))
+      # `x` and `y` will be fetched by the gradient computation, but not `loss`.
+      return loss
+
+    if isinstance(optimizer, optimizer_v2.OptimizerV2):
+      return optimizer.minimize(loss_fn, lambda: layer.trainable_variables)
+
+    # Callable loss.
+    return optimizer.minimize(loss_fn)
+
+  return model_fn, dataset_fn, batchnorm