From c608d46195cbc600c900c3d4fbf419ea63872827 Mon Sep 17 00:00:00 2001 From: Scott Zhu Date: Mon, 16 Mar 2020 23:10:19 -0700 Subject: [PATCH] Update tests under keras.mixed_precision to use combinations. Change all test_util.run_all_in_graph_and_eager_modes to combination. PiperOrigin-RevId: 301311255 Change-Id: Icd4f80a3e84ab06ef7d3423298918d172d1338e6 --- .../keras/mixed_precision/experimental/BUILD | 5 ++ .../experimental/autocast_variable_test.py | 29 +++++----- .../device_compatibility_check_test.py | 4 +- .../experimental/keras_test.py | 39 ++++++------- .../experimental/loss_scale_optimizer_test.py | 55 +++++++------------ .../experimental/policy_test.py | 8 ++- 6 files changed, 63 insertions(+), 77 deletions(-) diff --git a/tensorflow/python/keras/mixed_precision/experimental/BUILD b/tensorflow/python/keras/mixed_precision/experimental/BUILD index e364f346746..c672bd51cc7 100644 --- a/tensorflow/python/keras/mixed_precision/experimental/BUILD +++ b/tensorflow/python/keras/mixed_precision/experimental/BUILD @@ -63,6 +63,7 @@ py_test( "//tensorflow/python:client_testlib", "//tensorflow/python:platform_test", "//tensorflow/python/keras", + "//tensorflow/python/keras:combinations", "//tensorflow/python/keras/mixed_precision/experimental:loss_scale_optimizer", "//tensorflow/python/keras/optimizer_v2", ], @@ -86,6 +87,7 @@ cuda_py_test( deps = [ ":device_compatibility_check", "//tensorflow/python:client_testlib", + "//tensorflow/python/keras:combinations", ], ) @@ -137,6 +139,7 @@ py_test( "//tensorflow/python:platform_test", "//tensorflow/python/distribute:mirrored_strategy", "//tensorflow/python/eager:context", + "//tensorflow/python/keras:combinations", "//tensorflow/python/keras/optimizer_v2", "@absl_py//absl/testing:parameterized", ], @@ -177,6 +180,7 @@ cuda_py_test( "//tensorflow/python/distribute:mirrored_strategy", "//tensorflow/python/distribute:one_device_strategy", "//tensorflow/python/keras", + "//tensorflow/python/keras:combinations", ], ) @@ -219,6 +223,7 @@ cuda_py_test( "//tensorflow/python/distribute:central_storage_strategy", "//tensorflow/python/distribute:mirrored_strategy", "//tensorflow/python/keras", + "//tensorflow/python/keras:combinations", "@absl_py//absl/testing:parameterized", ], ) diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py index 204afd3913e..6ce3079788c 100644 --- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py +++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py @@ -30,7 +30,7 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import indexed_slices from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util +from tensorflow.python.keras import combinations from tensorflow.python.keras.mixed_precision.experimental import autocast_variable from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_v2 from tensorflow.python.ops import array_ops @@ -69,7 +69,7 @@ def get_var(val, dtype, name=None): return variables.VariableV1(val, use_resource=True, dtype=dtype, name=name) -@test_util.run_all_in_graph_and_eager_modes +@combinations.generate(combinations.combine(mode=['graph', 'eager'])) class AutoCastVariableTest(test.TestCase, parameterized.TestCase): @parameterized.named_parameters(*TESTCASES) @@ -156,7 +156,7 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase): def test_method_delegations(self, distribute): # Test AutoCastVariable correctly delegates Variable methods to the # underlying variable. - with get_distribute_scope(distribute): + with self.test_session(), get_distribute_scope(distribute): for read_dtype in (dtypes.float32, dtypes.float16): if distribute: # MirroredVariable.assign will (incorrectly) return a Mirrored value @@ -383,18 +383,19 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase): @parameterized.named_parameters(*TESTCASES) def test_checkpoint(self, distribute): - with get_distribute_scope(distribute): - x = get_var(1., dtypes.float32) - x = autocast_variable.create_autocast_variable(x) - self.evaluate(x.initializer) - self.evaluate(x.assign(123.)) + with self.test_session(): + with get_distribute_scope(distribute): + x = get_var(1., dtypes.float32) + x = autocast_variable.create_autocast_variable(x) + self.evaluate(x.initializer) + self.evaluate(x.assign(123.)) - checkpoint = trackable_utils.Checkpoint(x=x) - prefix = os.path.join(self.get_temp_dir(), 'ckpt') - save_path = checkpoint.save(prefix) - self.evaluate(x.assign(234.)) - checkpoint.restore(save_path).assert_consumed().run_restore_ops() - self.assertEqual(self.evaluate(x), 123.) + checkpoint = trackable_utils.Checkpoint(x=x) + prefix = os.path.join(self.get_temp_dir(), 'ckpt') + save_path = checkpoint.save(prefix) + self.evaluate(x.assign(234.)) + checkpoint.restore(save_path).assert_consumed().run_restore_ops() + self.assertEqual(self.evaluate(x), 123.) @parameterized.named_parameters(*TESTCASES) def test_invalid_wrapped_variable(self, distribute): diff --git a/tensorflow/python/keras/mixed_precision/experimental/device_compatibility_check_test.py b/tensorflow/python/keras/mixed_precision/experimental/device_compatibility_check_test.py index c3315caab88..a88594ae346 100644 --- a/tensorflow/python/keras/mixed_precision/experimental/device_compatibility_check_test.py +++ b/tensorflow/python/keras/mixed_precision/experimental/device_compatibility_check_test.py @@ -21,7 +21,7 @@ from __future__ import print_function import re from tensorflow.core.framework import device_attributes_pb2 -from tensorflow.python.framework import test_util +from tensorflow.python.keras import combinations from tensorflow.python.keras.mixed_precision.experimental import device_compatibility_check from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging @@ -42,7 +42,7 @@ def _get_device_attrs(device_type, device_name=None, cc_major=None, device_type='GPU', physical_device_desc=physical_device_desc) -@test_util.run_all_in_graph_and_eager_modes +@combinations.generate(combinations.combine(mode=['graph', 'eager'])) class DeviceCompatibilityCheckTest(test.TestCase): def _test_compat_check(self, device_attr_list, should_warn, expected_regex, diff --git a/tensorflow/python/keras/mixed_precision/experimental/keras_test.py b/tensorflow/python/keras/mixed_precision/experimental/keras_test.py index a0de45d37fb..c46ee9123ce 100644 --- a/tensorflow/python/keras/mixed_precision/experimental/keras_test.py +++ b/tensorflow/python/keras/mixed_precision/experimental/keras_test.py @@ -31,8 +31,8 @@ from tensorflow.python.eager import context from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes -from tensorflow.python.framework import test_util from tensorflow.python.keras import backend +from tensorflow.python.keras import combinations from tensorflow.python.keras import keras_parameterized from tensorflow.python.keras import layers from tensorflow.python.keras import models @@ -119,11 +119,11 @@ TESTCASES = ({ }) +@combinations.generate(combinations.combine(mode=['graph', 'eager'])) class KerasLayerTest(keras_parameterized.TestCase): """Test mixed precision with Keras layers.""" @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_mixed_policies_(self, strategy_fn): for dtype in 'float16', 'bfloat16': x = constant_op.constant([1.]) @@ -142,7 +142,6 @@ class KerasLayerTest(keras_parameterized.TestCase): self.evaluate(variables.global_variables_initializer()) self.assertEqual(self.evaluate(y), 1.) - @test_util.run_in_graph_and_eager_modes def test_layer_with_int_variable(self): class LayerWithIntVar(base_layer.Layer): @@ -159,7 +158,6 @@ class KerasLayerTest(keras_parameterized.TestCase): self.assertEqual(layer(x).dtype, 'int32') @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_layer_with_non_autocast_variable(self, strategy_fn): x = constant_op.constant([1.]) with strategy_fn().scope(): @@ -172,7 +170,6 @@ class KerasLayerTest(keras_parameterized.TestCase): self.assertEqual(self.evaluate(y), 1.) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_layer_calling_tf_function(self, strategy_fn): x = constant_op.constant([1.]) with strategy_fn().scope(): @@ -185,7 +182,6 @@ class KerasLayerTest(keras_parameterized.TestCase): self.assertEqual(self.evaluate(y), 1.) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_layer_regularizer_runs_in_var_dtype(self, strategy_fn): x = constant_op.constant([1.]) with strategy_fn().scope(): @@ -211,7 +207,6 @@ class KerasLayerTest(keras_parameterized.TestCase): self.assertEqual(self.evaluate(regularizer_loss), 1.) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_passing_policy_to_layer(self, strategy_fn): x = constant_op.constant([1.], dtype=dtypes.float16) with strategy_fn().scope(): @@ -230,7 +225,6 @@ class KerasLayerTest(keras_parameterized.TestCase): self.assertEqual(layer(x).dtype, dtypes.float64) self.assertEqual(layer.v.dtype, dtypes.float64) - @test_util.run_in_graph_and_eager_modes def test_error_passing_policy_string_to_layer(self): with self.assertRaisesRegexp( TypeError, "Cannot convert value 'mixed_float16' to a " @@ -240,7 +234,6 @@ class KerasLayerTest(keras_parameterized.TestCase): mp_test_util.MultiplyLayer(dtype='mixed_float16') @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_gradient(self, strategy_fn): x = constant_op.constant([1.]) with strategy_fn().scope() as strategy: @@ -308,17 +301,19 @@ class KerasLayerTest(keras_parameterized.TestCase): self.assertEqual(self.evaluate(layer(x)), 100.) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_checkpointing_layer_weights(self, strategy_fn): - self._test_checkpointing_layer_weights( - strategy_fn, mixed_prec_when_saving=True, mixed_prec_when_loading=True) - self._test_checkpointing_layer_weights( - strategy_fn, mixed_prec_when_saving=True, mixed_prec_when_loading=False) - self._test_checkpointing_layer_weights( - strategy_fn, mixed_prec_when_saving=False, mixed_prec_when_loading=True) + with self.test_session(): + self._test_checkpointing_layer_weights( + strategy_fn, mixed_prec_when_saving=True, + mixed_prec_when_loading=True) + self._test_checkpointing_layer_weights( + strategy_fn, mixed_prec_when_saving=True, + mixed_prec_when_loading=False) + self._test_checkpointing_layer_weights( + strategy_fn, mixed_prec_when_saving=False, + mixed_prec_when_loading=True) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def test_config(self, strategy_fn): x = constant_op.constant([1.], dtype=dtypes.float16) with strategy_fn().scope(): @@ -397,7 +392,6 @@ class KerasLayerTest(keras_parameterized.TestCase): self.assertEqual(layer(x).dtype, 'float16') self.assertEqual(layer.v.dtype, 'float16') - @test_util.run_in_graph_and_eager_modes def test_delete_variable(self): layer = base_layer.Layer(dtype=policy.Policy('mixed_float16')) layer.x = layer.add_weight('x') @@ -405,7 +399,6 @@ class KerasLayerTest(keras_parameterized.TestCase): del layer.x self.assertEqual(layer.trainable_weights, []) - @test_util.run_in_graph_and_eager_modes def test_build_and_call_layer_in_function(self): layer = mp_test_util.MultiplyLayer(dtype=policy.Policy('mixed_float16')) @def_function.function @@ -833,7 +826,7 @@ class KerasModelTest(keras_parameterized.TestCase): model.fit(dataset) self.assertEqual(backend.eval(layer.v), -3) - @test_util.run_in_graph_and_eager_modes + @combinations.generate(combinations.combine(mode=['graph', 'eager'])) def test_loss_scale_optimizer_overrides_policy_loss_scale(self): with policy.policy_scope(policy.Policy('float32', loss_scale=10.)): opt = gradient_descent.SGD(1.) @@ -844,7 +837,7 @@ class KerasModelTest(keras_parameterized.TestCase): model.compile(opt, loss='mse') self.assertEqual(self.evaluate(model.optimizer.loss_scale()), 5.) - @test_util.run_in_graph_and_eager_modes + @combinations.generate(combinations.combine(mode=['graph', 'eager'])) def test_pass_invalid_optimizer_with_loss_scaling(self): with policy.policy_scope(policy.Policy('float32', loss_scale=10.)): x = layers.Input(shape=(1,)) @@ -857,7 +850,7 @@ class KerasModelTest(keras_parameterized.TestCase): with self.assertRaisesRegexp(ValueError, error_msg): model.compile(optimizers.SGD(1.), 'mse') - @test_util.run_in_graph_and_eager_modes + @combinations.generate(combinations.combine(mode=['graph', 'eager'])) def test_functional_model_loss_dtype(self): with policy.policy_scope('float16'): x = layers.Input(shape=(1,)) @@ -867,6 +860,7 @@ class KerasModelTest(keras_parameterized.TestCase): # The loss should not be casted to the policy's dtype. self.assertEqual(model.losses[0].dtype, 'float32') + @keras_parameterized.run_all_keras_modes @parameterized.named_parameters( { 'testcase_name': 'base', @@ -883,7 +877,6 @@ class KerasModelTest(keras_parameterized.TestCase): 'strategy_fn': create_mirrored_strategy, 'h5': True, }) - @test_util.run_in_graph_and_eager_modes def test_save_weights_with_autocast_vars(self, strategy_fn, h5=False): with strategy_fn().scope(): with policy.policy_scope('mixed_float16'): diff --git a/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer_test.py b/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer_test.py index 6a6b3cb8dc2..8d13c6a37a3 100644 --- a/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer_test.py +++ b/tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer_test.py @@ -28,6 +28,7 @@ from tensorflow.python.distribute import mirrored_strategy from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.framework import test_util +from tensorflow.python.keras import combinations from tensorflow.python.keras import optimizers from tensorflow.python.keras.mixed_precision.experimental import loss_scale_optimizer from tensorflow.python.keras.mixed_precision.experimental import test_util as mp_test_util @@ -68,6 +69,7 @@ TESTCASES = ({ @test_util.with_control_flow_v2 +@combinations.generate(combinations.combine(mode=['graph', 'eager'])) class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): def _run_if_in_graph_mode(self, val): @@ -85,7 +87,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): return lambda: opt.minimize(loss, var_list=[var]) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def testFixedLossScaleAppliedToLossWithMinimize(self, strategy_fn): with strategy_fn().scope() as strategy: var = variables.Variable([5.0]) @@ -119,7 +120,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): # mp_test_util.create_identity_with_grad_check_fn added an assertion op. self.evaluate(run_op) - @test_util.run_in_graph_and_eager_modes def testGetScaledLoss(self): opt = gradient_descent.SGD(2.0) opt = loss_scale_optimizer.LossScaleOptimizer(opt, loss_scale=2.) @@ -130,7 +130,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertEqual(10., self.evaluate(opt.get_scaled_loss(loss))) self.assertEqual(10., self.evaluate(opt.get_scaled_loss(lambda: loss)())) - @test_util.run_in_graph_and_eager_modes def testGetUnscaledGradients(self): opt = gradient_descent.SGD(2.0) opt = loss_scale_optimizer.LossScaleOptimizer(opt, loss_scale=2) @@ -142,7 +141,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): grads = [self.evaluate(g) if g is not None else g for g in grads] self.assertEqual([1.5, None, -2.], grads) - @test_util.run_in_graph_and_eager_modes def testGetUnscaledSparseGradients(self): opt = gradient_descent.SGD(2.0) opt = loss_scale_optimizer.LossScaleOptimizer(opt, loss_scale=2) @@ -156,7 +154,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.evaluate(sparse_grad.values)) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def testDynamicLossScale(self, strategy_fn): strategy = strategy_fn() learning_rate = 2. @@ -190,7 +187,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertAllClose([1.], self.evaluate(var)) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def testDynamicUpdate(self, strategy_fn): with strategy_fn().scope() as strategy: var = variables.Variable([1.0, 2.0]) @@ -221,7 +217,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertEqual(2., self.evaluate(opt.loss_scale())) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def testDynamicLossScaleWithFloat16Loss(self, strategy_fn): strategy = strategy_fn() learning_rate = 2. @@ -243,7 +238,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertAllClose([3.], self.evaluate(var)) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def testDynamicLossScaleWithSlots(self, strategy_fn): strategy_obj = strategy_fn() if (isinstance(strategy_obj, mirrored_strategy.MirroredStrategy) and @@ -283,7 +277,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertEqual(opt.get_slot_names(), ['momentum']) - @test_util.run_in_graph_and_eager_modes def testIterations(self): opt = gradient_descent.SGD(2.0) lso = loss_scale_optimizer.LossScaleOptimizer(opt, loss_scale=10.) @@ -291,27 +284,26 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertEqual(lso.iterations, 7) self.assertEqual(opt.iterations, 7) - @test_util.run_in_graph_and_eager_modes def testWeightMethods(self): - var = variables.Variable([1.0]) - opt = gradient_descent.SGD(1.0) - initial_loss_scale = 2. - loss_scale = loss_scale_module.DynamicLossScale( - initial_loss_scale=initial_loss_scale, increment_period=1, - multiplier=4) - opt = loss_scale_optimizer.LossScaleOptimizer(opt, loss_scale) - run_op = opt.minimize(lambda: var * 2, [var]) - self.evaluate(variables.global_variables_initializer()) - self._run_if_in_graph_mode(run_op) + with self.test_session(): + var = variables.Variable([1.0]) + opt = gradient_descent.SGD(1.0) + initial_loss_scale = 2. + loss_scale = loss_scale_module.DynamicLossScale( + initial_loss_scale=initial_loss_scale, increment_period=1, + multiplier=4) + opt = loss_scale_optimizer.LossScaleOptimizer(opt, loss_scale) + run_op = opt.minimize(lambda: var * 2, [var]) + self.evaluate(variables.global_variables_initializer()) + self._run_if_in_graph_mode(run_op) - self.assertLen(opt.weights, 1) # The 'iterations' weight - self.assertEqual(self.evaluate(opt.weights[0]), 1) - self.assertEqual(opt.get_weights()[0], 1) - self.assertEqual(self.evaluate(opt.variables()[0]), 1) - opt.set_weights([np.array(2.)]) - self.assertEqual(self.evaluate(opt.variables()[0]), 2) + self.assertLen(opt.weights, 1) # The 'iterations' weight + self.assertEqual(self.evaluate(opt.weights[0]), 1) + self.assertEqual(opt.get_weights()[0], 1) + self.assertEqual(self.evaluate(opt.variables()[0]), 1) + opt.set_weights([np.array(2.)]) + self.assertEqual(self.evaluate(opt.variables()[0]), 2) - @test_util.run_in_graph_and_eager_modes def testSlotMethodErrors(self): opt = gradient_descent.SGD(1.0, momentum=1.0) opt = loss_scale_optimizer.LossScaleOptimizer(opt, 'dynamic') @@ -332,9 +324,8 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): loss_scale_optimizer.LossScaleOptimizer(opt, None) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def testGettingAndSettingLearningRate(self, strategy_fn): - with strategy_fn().scope() as strategy: + with self.test_session(), strategy_fn().scope() as strategy: var = variables.Variable([5.0]) opt = adam.Adam(learning_rate=1.0) loss = lambda: var * 2.0 @@ -357,7 +348,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): with self.assertRaises(AttributeError): opt.not_an_attr += 3 - @test_util.run_in_graph_and_eager_modes def testArbitraryAttributesNotExposed(self): opt = adam.Adam(learning_rate=1.0) # Test that Adam has attributes 'epsilon' and 'beta1' @@ -375,7 +365,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): "'LossScaleOptimizer' object has no attribute 'beta_1'"): opt.beta_1 # pylint: disable=pointless-statement - @test_util.run_in_graph_and_eager_modes def testApplyGradientsGetsUnwrappedTensors(self): # Tests that gradients passed to apply_gradients are not wrapped in a # DistributionStrategy wrapper, such as PerReplica, but instead are raw @@ -403,7 +392,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): strategy.experimental_run(run_fn) @parameterized.named_parameters(*TESTCASES) - @test_util.run_in_graph_and_eager_modes def testCheckpoint(self, strategy_fn): strategy = strategy_fn() if (isinstance(strategy, mirrored_strategy.MirroredStrategy) and @@ -447,7 +435,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertEqual(self.evaluate(loss_scale._num_good_steps), 1) self.assertAlmostEqual(self.evaluate(slot_var).item(), slot_value) - @test_util.run_in_graph_and_eager_modes def testGetConfig(self): opt = gradient_descent.SGD(2., momentum=0.5) loss_scale = loss_scale_module.DynamicLossScale( @@ -466,7 +453,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertEqual(opt.loss_scale.increment_period, 3.) self.assertEqual(opt.loss_scale.multiplier, 4.) - @test_util.run_in_graph_and_eager_modes def testSerializationWithBuiltInOptimizer(self): opt = gradient_descent.SGD(2., momentum=0.5) loss_scale = loss_scale_module.DynamicLossScale( @@ -485,7 +471,6 @@ class LossScaleOptimizerTest(test.TestCase, parameterized.TestCase): self.assertEqual(opt.loss_scale.increment_period, 3.) self.assertEqual(opt.loss_scale.multiplier, 4.) - @test_util.run_in_graph_and_eager_modes def testSerializationWithCustomOptimizer(self): class MySGD(gradient_descent.SGD): diff --git a/tensorflow/python/keras/mixed_precision/experimental/policy_test.py b/tensorflow/python/keras/mixed_precision/experimental/policy_test.py index 2ea705852d6..b345039b406 100644 --- a/tensorflow/python/keras/mixed_precision/experimental/policy_test.py +++ b/tensorflow/python/keras/mixed_precision/experimental/policy_test.py @@ -18,10 +18,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized + from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util +from tensorflow.python.keras import combinations from tensorflow.python.keras import testing_utils from tensorflow.python.keras.engine import base_layer_utils from tensorflow.python.keras.mixed_precision.experimental import device_compatibility_check @@ -33,8 +35,8 @@ from tensorflow.python.training.experimental import loss_scale as loss_scale_mod from tensorflow.python.training.experimental import mixed_precision -@test_util.run_all_in_graph_and_eager_modes -class PolicyTest(test.TestCase): +@combinations.generate(combinations.combine(mode=['graph', 'eager'])) +class PolicyTest(test.TestCase, parameterized.TestCase): """Tests Policies.""" @testing_utils.enable_v2_dtype_behavior