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
This commit is contained in:
Scott Zhu 2020-03-16 23:10:19 -07:00 committed by TensorFlower Gardener
parent e473416af2
commit c608d46195
6 changed files with 63 additions and 77 deletions

View File

@ -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",
],
)

View File

@ -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):

View File

@ -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,

View File

@ -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'):

View File

@ -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):

View File

@ -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