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:
parent
e473416af2
commit
c608d46195
@ -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",
|
||||
],
|
||||
)
|
||||
|
@ -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):
|
||||
|
@ -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,
|
||||
|
@ -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'):
|
||||
|
@ -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):
|
||||
|
||||
|
@ -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
|
||||
|
Loading…
x
Reference in New Issue
Block a user