def test_error_if_policy_is_set(self):
with policy.policy_scope('infer_float32_vars'):
with self.assertRaisesRegexp(
ValueError, 'a keras mixed precision Policy has been set'):
mixed_precision.enable_mixed_precision_graph_rewrite(
gradient_descent_v2.SGD(1.0))
# Test no error is thrown when the policy is current the default.
mixed_precision.enable_mixed_precision_graph_rewrite(
gradient_descent_v2.SGD(1.0))
def test_error_if_graph_rewrite_enabled(self):
try:
mixed_precision.enable_mixed_precision_graph_rewrite(
gradient_descent.SGD(1.))
with self.assertRaisesRegexp(
ValueError, 'the mixed precision graph rewrite has already been '
'enabled'):
mp_policy.set_policy('infer_float32_vars')
finally:
mixed_precision.disable_mixed_precision_graph_rewrite()
def test_wrap_optimizer(self):
opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
opt = mixed_precision.enable_mixed_precision_graph_rewrite(opt, 123.)
self.assertIsInstance(
opt, loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer)
self.assertEqual(self.evaluate(opt._loss_scale()), 123.)
opt = gradient_descent_v2.SGD(1.0)
opt = mixed_precision.enable_mixed_precision_graph_rewrite(opt, 123.)
self.assertIsInstance(opt, loss_scale_optimizer_v2.LossScaleOptimizer)
self.assertEqual(self.evaluate(opt._loss_scale()), 123.)
def test_error_if_graph_rewrite_enabled(self):
try:
mixed_precision.enable_mixed_precision_graph_rewrite(
gradient_descent.SGD(1.))
with self.assertRaisesRegex(
ValueError, 'cannot be set to "mixed_float16", .* the mixed '
'precision graph rewrite has already been enabled'):
mp_policy.set_policy('mixed_float16')
with mp_policy.policy_scope('float64'):
pass # Non-mixed policies are allowed
finally:
mixed_precision.disable_mixed_precision_graph_rewrite()
def test_warn_if_session_already_exists(self, mock_warn):
# Set this to False, so Sessions created in previous tests do not trigger
# the warning.
mixed_precision_global_state.non_mixed_precision_session_created = False
with session.Session():
mixed_precision.enable_mixed_precision_graph_rewrite(
gradient_descent_v2.SGD(1.0))
mock_warn.assert_any_call(
'You already have existing Sessions that do not use mixed precision. '
'enable_mixed_precision_graph_rewrite() will not affect these '
'Sessions.')
def test_do_not_warn_if_session_does_not_already_exist(self, mock_warn):
# Set this to False, so Sessions created in previous tests do not trigger
# the warning.
mixed_precision_global_state.non_mixed_precision_session_created = False
mixed_precision.enable_mixed_precision_graph_rewrite(
gradient_descent_v2.SGD(1.0))
with session.Session():
# Make sure the "You already have existing Sessions" warning was not
# issued, since the Session was only created after
# enable_mixed_precision_graph_rewrite.
for call_arg in mock_warn.call_args_list:
msg = call_arg[0][0]
self.assertNotIn('You already have existing Sessions that do not use '
'mixed precision', msg)
def test_optimizer_errors(self):
opt = 1
expected_regex = (
'"opt" must be an instance of a tf.train.Optimizer or '
'a tf.keras.optimizers.Optimizer, but got')
with self.assertRaisesRegexp(ValueError, expected_regex):
mixed_precision.enable_mixed_precision_graph_rewrite(opt)
self.assertFalse(config.get_optimizer_experimental_options().get(
'auto_mixed_precision', False))
opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
opt = loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer(
opt, 'dynamic')
with self.assertRaisesRegexp(
ValueError, '"opt" must not already be an instance of a '
'MixedPrecisionLossScaleOptimizer.'):
mixed_precision.enable_mixed_precision_graph_rewrite(opt)
self.assertFalse(config.get_optimizer_experimental_options().get(
'auto_mixed_precision', False))
opt = gradient_descent_v2.SGD(1.0)
opt = loss_scale_optimizer_v2.LossScaleOptimizer(opt, 'dynamic')
with self.assertRaisesRegexp(
ValueError, '"opt" must not already be an instance of a '
'LossScaleOptimizer.'):
mixed_precision.enable_mixed_precision_graph_rewrite(opt)
self.assertFalse(config.get_optimizer_experimental_options().get(
'auto_mixed_precision', False))
def test_grappler_pass_enabled(self):
opt = gradient_descent_v2.SGD(1.0)
mixed_precision.enable_mixed_precision_graph_rewrite(opt, 123.)
var = variables.Variable([[1.0]])
def overflow_in_float16():
out = var * 2**10
out = math_ops.matmul(out, out)
return array_ops.reshape(out, ())
if context.executing_eagerly():
f = def_function.function(overflow_in_float16)
self.assertEqual(f().numpy(), float('Inf'))
# Outside a def_function.function, the grappler pass will not be applied.
self.assertAlmostEqual(overflow_in_float16().numpy(), 2**20)
# Test disabling mixed precision.
mixed_precision.disable_mixed_precision_graph_rewrite()
self.assertEqual(f().numpy(), 2**20)
else:
with session.Session() as sess:
out = overflow_in_float16()
sess.run(var.initializer)
self.assertEqual(sess.run(out), float('Inf'))
# Test Session will enable the auto_mixed_precision grappler pass in a
# ConfigProto passed by the user
with session.Session(config=config_pb2.ConfigProto()) as sess:
out = overflow_in_float16()
sess.run(var.initializer)
self.assertEqual(sess.run(out), float('Inf'))
# Test disabling mixed precision.
mixed_precision.disable_mixed_precision_graph_rewrite()
with session.Session() as sess:
out = overflow_in_float16()
sess.run(var.initializer)
self.assertAlmostEqual(sess.run(out), 2**20)