def all_strategy_combinations_with_eager_and_graph_modes():
return (
combinations.combine(distribution=strategy_combinations.all_strategies,
mode=['graph', 'eager']) +
combinations.combine(distribution=strategy_combinations.
multi_worker_mirrored_strategies,
mode='eager'))
class TestDistributionStrategyDnnCorrectness(test.TestCase,
parameterized.TestCase):
"""Test custom training loop correctness with a simple DNN model."""
def setUp(self):
super(TestDistributionStrategyDnnCorrectness, self).setUp()
v2_compat.enable_v2_behavior()
np.random.seed(_RANDOM_SEED)
random_seed.set_random_seed(_RANDOM_SEED)
@ds_combinations.generate(
combinations.combine(
distribution=strategy_combinations.all_strategies,
optimizer_fn=optimizer_combinations.optimizers_v2,
mode=['eager'],
iteration_type=['iterator', 'dataset'],
inside_func=[False, True],
sync_batchnorm=[True, False]) +
combinations.combine(
distribution=strategy_combinations.multiworker_strategies,
optimizer_fn=[
optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
optimizer_combinations.adagrad_optimizer_keras_v2_fn
],
mode=['eager'],
iteration_type=['iterator', 'dataset'],
inside_func=[False, True],
sync_batchnorm=[True, False]
))
def test_dnn_correctness_minus_tpus(self, distribution, optimizer_fn,
iteration_type, inside_func,
sync_batchnorm):
# TODO(anjs): Identify why this particular V1 optimizer needs a higher tol.
if 'FtrlV1' in optimizer_fn._name and 'TPU' in type(distribution).__name__:
self.skipTest('Reduced tolerance of the order of 1e-1 required.')
if ('CollectiveAllReduce' in type(distribution).__name__ and
test_util.is_xla_enabled()):
self.skipTest('XLA tests fail with MWMS.')
self.dnn_correctness(distribution, optimizer_fn, iteration_type,
inside_func, sync_batchnorm)
def dnn_correctness(self, distribution, optimizer_fn, iteration_type,
inside_func, sync_batchnorm=None):
model = get_model(sync_batchnorm)
initial_weights = model.get_weights()
dataset = get_data()
if inside_func:
iteration_func = iteration_inside_func
else:
iteration_func = iteration_outside_func
wts_with_ds, loss_with_ds, acc_with_ds = iteration_func(
initial_weights, dataset, optimizer_fn, iteration_type,
strategy=distribution, sync_batchnorm=sync_batchnorm)
wts, loss, acc = iteration_func(initial_weights, dataset, optimizer_fn,
iteration_type,
sync_batchnorm=sync_batchnorm)
self.assertAllClose(wts, wts_with_ds, atol=1e-3, rtol=1e-3)
self.assertAllClose(loss, loss_with_ds, atol=1e-3, rtol=1e-3)
self.assertAllClose(acc, acc_with_ds, atol=1e-3, rtol=1e-3)
class DistributedCollectiveAllReduceStrategyTestWithChief(
CollectiveAllReduceStrategyTestBase, parameterized.TestCase):
@classmethod
def setUpClass(cls):
"""Create a local cluster with 3 workers and 1 chief."""
cls._cluster_spec = multi_worker_test_base.create_in_process_cluster(
num_workers=3, num_ps=0, has_chief=True)
@ds_combinations.generate(
combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
def testComplexModel(self, required_gpus):
self._run_between_graph_clients(self._test_complex_model,
self._cluster_spec,
num_gpus=required_gpus)
@ds_combinations.generate(
combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
@testing_utils.enable_v2_dtype_behavior
def testMixedPrecision(self, required_gpus):
if test_util.is_xla_enabled():
return # Test gets NaNs with XLA
with policy.policy_scope('mixed_float16'):
self._run_between_graph_clients(self._test_mixed_precision,
self._cluster_spec,
num_gpus=required_gpus)
def test_times_variable_arguments(self):
c1 = combinations.combine(mode=["graph", "eager"])
c2 = combinations.combine(optimizer=["adam", "gd"])
c3 = combinations.combine(distribution=["d1", "d2"])
c4 = combinations.times(c3, c1, c2)
self.assertEqual([
OrderedDict([("distribution", "d1"), ("mode", "graph"),
("optimizer", "adam")]),
OrderedDict([("distribution", "d1"), ("mode", "graph"),
("optimizer", "gd")]),
OrderedDict([("distribution", "d1"), ("mode", "eager"),
("optimizer", "adam")]),
OrderedDict([("distribution", "d1"), ("mode", "eager"),
("optimizer", "gd")]),
OrderedDict([("distribution", "d2"), ("mode", "graph"),
("optimizer", "adam")]),
OrderedDict([("distribution", "d2"), ("mode", "graph"),
("optimizer", "gd")]),
OrderedDict([("distribution", "d2"), ("mode", "eager"),
("optimizer", "adam")]),
OrderedDict([("distribution", "d2"), ("mode", "eager"),
("optimizer", "gd")])
], c4)
self.assertEqual(
combinations.combine(mode=["graph", "eager"],
optimizer=["adam", "gd"],
distribution=["d1", "d2"]), c4)
class SavedModelKerasModelTest(test_base.TestSavedModelBase):
def setUp(self):
self._root_dir = 'saved_model_save_load'
super(SavedModelKerasModelTest, self).setUp()
def _save_model(self, model, saved_dir):
saved_model.save(model, saved_dir)
def _load_and_run_model(self,
distribution,
saved_dir,
predict_dataset,
output_name='output_1'):
return test_base.load_and_run_with_saved_model_api(distribution, saved_dir,
predict_dataset,
output_name)
@ds_combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution):
self.run_test_save_no_strategy_restore_strategy(
model_and_input, distribution)
@ds_combinations.generate(
combinations.times(test_base.simple_models_with_strategies(),
combinations.combine(save_in_scope=[True, False])))
def test_save_strategy_restore_no_strategy(self, model_and_input,
distribution, save_in_scope):
self.run_test_save_strategy_restore_no_strategy(
model_and_input, distribution, save_in_scope)
@ds_combinations.generate(
combinations.times(test_base.simple_models_with_strategy_pairs(),
combinations.combine(save_in_scope=[True, False])))
def test_save_strategy_restore_strategy(self, model_and_input,
distribution_for_saving,
distribution_for_restoring,
save_in_scope):
self.run_test_save_strategy_restore_strategy(model_and_input,
distribution_for_saving,
distribution_for_restoring,
save_in_scope)
@ds_combinations.generate(
combinations.times(test_base.simple_models_with_strategies(),
combinations.combine(save_in_scope=[True, False])))
def test_no_variable_device_placement(self, model_and_input, distribution,
save_in_scope):
saved_dir = self.run_test_save_strategy(model_and_input, distribution,
save_in_scope)
func = saved_model.load(saved_dir)
concrete_function = func.signatures[test_base._DEFAULT_FUNCTION_KEY]
for f in concrete_function.graph.as_graph_def().library.function:
for n in f.node_def:
if n.op == 'ReadVariableOp':
self.assertEmpty(n.device)
def test_combinations_for_embedding_model():
# TODO(sourabhbajaj): Enable tests for eager mode
eager_mode_strategies = [
s for s in strategies_for_embedding_models() if not s.required_tpu
]
return (combinations.times(
combinations.combine(distribution=strategies_for_embedding_models()),
(graph_mode_test_configuration())) + combinations.times(
combinations.combine(distribution=eager_mode_strategies),
(eager_mode_test_configuration())))
def test_add(self):
self.assertEqual([{
"a": 1
}, {
"a": 2
}, {
"b": 2
}, {
"b": 3
}],
combinations.combine(a=[1, 2]) +
combinations.combine(b=[2, 3]))
class TestDistributionStrategySaveLoadWeights(test.TestCase,
parameterized.TestCase):
@ds_combinations.generate(
combinations.times(
keras_test_lib.all_strategy_combinations_minus_default(),
combinations.combine(optimizer=optimizer_combinations.
rmsprop_optimizer_keras_v2_fn)))
def test_save_load_h5(self, distribution, optimizer):
with self.cached_session():
dataset = keras_test_lib.get_dataset(distribution)
with distribution.scope():
model = keras_test_lib.get_model()
model.compile(optimizer(), 'mse')
model.fit(dataset, epochs=1, steps_per_epoch=1)
weights_file = tempfile.mktemp('.h5')
model.save_weights(weights_file)
model_2 = keras_test_lib.get_model()
model_2.compile(optimizer(), 'mse')
model_2.load_weights(weights_file)
model_2.predict(
keras_test_lib.get_predict_dataset(distribution), steps=2)
model_2.fit(dataset, epochs=1, steps_per_epoch=1)
@ds_combinations.generate(
combinations.times(
keras_test_lib.all_strategy_combinations_minus_default(),
combinations.combine(optimizer=optimizer_combinations.
rmsprop_optimizer_keras_v2_fn)))
def test_save_load_trackable(self, distribution, optimizer):
# TODO(b/123533246): Enable the test for TPU once bug is fixed
if (isinstance(distribution,
(tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1))
and distribution.extended.steps_per_run > 1):
self.skipTest(
'MultiStep TPU Strategy deadlocks with optimizer restore.')
with self.cached_session():
dataset = keras_test_lib.get_dataset(distribution)
with distribution.scope():
model = keras_test_lib.get_model()
model.compile(optimizer(), 'mse')
model.fit(dataset, epochs=1, steps_per_epoch=1)
weights_file = tempfile.mktemp()
model.save_weights(weights_file)
model_2 = keras_test_lib.get_model()
model_2.compile(optimizer(), 'mse')
model_2.load_weights(weights_file)
model_2.predict(
keras_test_lib.get_predict_dataset(distribution), steps=2)
model_2.fit(dataset, epochs=1, steps_per_epoch=1)
def test_arguments_sorted(self):
self.assertEqual([
OrderedDict([("aa", 1), ("ab", 2)]),
OrderedDict([("aa", 1), ("ab", 3)]),
OrderedDict([("aa", 2), ("ab", 2)]),
OrderedDict([("aa", 2), ("ab", 3)])
], combinations.combine(ab=[2, 3], aa=[1, 2]))
class ModelCheckpointTest(test_base.IndependentWorkerTestBase,
parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(mode=['graph'],
required_gpus=[0, 1],
file_format=['h5', 'tf'],
save_weights_only=[True, False]))
def testCheckpointExists(self, file_format, save_weights_only):
train_ds, _ = multi_worker_testing_utils.mnist_synthetic_dataset(64, 2)
model = multi_worker_testing_utils.get_mnist_model((28, 28, 1))
saving_dir = self.get_temp_dir()
saving_filepath = os.path.join(saving_dir, 'checkpoint.' + file_format)
callbacks_list = [
callbacks.ModelCheckpoint(filepath=saving_filepath,
save_weights_only=save_weights_only)
]
self.assertFalse(file_io.file_exists_v2(saving_filepath))
try:
model.fit(x=train_ds,
epochs=2,
steps_per_epoch=2,
callbacks=callbacks_list)
except NotFoundError as e:
if 'Failed to create a NewWriteableFile' in e.message:
self.skipTest(
'b/138941852, path not found error in Windows py35.')
tf_saved_model_exists = file_io.file_exists_v2(saving_filepath)
tf_weights_only_checkpoint_exists = file_io.file_exists_v2(
saving_filepath + '.index')
self.assertTrue(tf_saved_model_exists
or tf_weights_only_checkpoint_exists)
def test_graph_and_eager_modes(test_class_or_method=None):
"""Decorator for generating graph and eager mode tests from a single test.
Must be applied to subclasses of `parameterized.TestCase` (from
absl/testing), or a method of such a subclass.
When applied to a test method, this decorator results in the replacement of
that method with a two new test methods, one executed in graph mode and the
other in eager mode.
When applied to a test class, all the methods in the class are affected.
Args:
test_class_or_method: the `TestCase` class or method to decorate.
Returns:
decorator: A generated TF `test_combinations` decorator, or if
`test_class_or_method` is not `None`, the generated decorator applied to
that function.
"""
decorator = test_combinations.generate(
test_combinations.combine(mode=['graph', 'eager']),
test_combinations=[combinations.EagerGraphCombination()])
if test_class_or_method:
return decorator(test_class_or_method)
return decorator
class KerasMultiWorkerTestIndependentWorker(test_base.IndependentWorkerTestBase,
parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(
mode=['graph'],
strategy_cls=[
collective_strategy.CollectiveAllReduceStrategy,
],
required_gpus=[0, 1]))
def testSimpleModelIndependentWorkerSync(self, strategy_cls):
num_workers = 2
num_epoch = 2
cluster_spec = test_base.create_cluster_spec(num_workers=num_workers)
self._barrier = dc._Barrier(2)
# The verification callback will be shared by multiple threads.
verification_callback = MultiWorkerVerificationCallback(
num_epoch=num_epoch, num_worker=num_workers)
def _independent_worker_fn(*args, **kwargs): # pylint: disable=unused-argument
"""Simulates an Independent Worker inside of a thread."""
with test.mock.patch.object(dc, '_run_std_server',
self._make_mock_run_std_server()):
strategy = strategy_cls()
verification_callback.is_between_graph = \
strategy.extended.experimental_between_graph
batch_size = 64
steps = 2
train_ds, _ = multi_worker_testing_utils.mnist_synthetic_dataset(
batch_size, steps)
with strategy.scope():
model = multi_worker_testing_utils.get_mnist_model((28, 28, 1))
orig_loss, _ = model.evaluate(train_ds, steps=steps)
callbacks_for_fit = nest.flatten(
kwargs.get('verification_callback', []))
history = model.fit(
x=train_ds,
epochs=num_epoch,
steps_per_epoch=steps,
callbacks=callbacks_for_fit)
self.assertIsInstance(history, keras.callbacks.History)
trained_loss, _ = model.evaluate(train_ds, steps=steps)
self.assertLess(trained_loss, orig_loss)
threads = self.run_multiple_tasks_in_threads(
_independent_worker_fn,
cluster_spec,
verification_callback=verification_callback)
threads_to_join = []
strategy = strategy_cls()
if strategy.extended.experimental_between_graph:
for ts in threads.values():
threads_to_join.extend(ts)
else:
threads_to_join = [threads['worker'][0]]
self.join_independent_workers(threads_to_join)
verification_callback.verify(self)
class KPLMultiWorkerTest(test.TestCase, parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(
mode=['eager'],
use_adapt=[False], # TODO(b/180742437): Add tests for using adapt.
strategy=[
strategy_combinations.multi_worker_mirrored_2x1_gpu,
# TODO(b/183956672): Re-enable
# strategy_combinations.multi_worker_mirrored_2x2_gpu,
]))
def testTrainAndServeWithKPL(self, use_adapt, strategy):
test_utils_obj = kpl_test_utils.DistributeKplTestUtils()
with strategy.scope():
feature_mapper, label_mapper = test_utils_obj.define_kpls_for_training(
use_adapt)
model = test_utils_obj.define_model()
optimizer = rmsprop.RMSprop(learning_rate=0.1)
accuracy = keras.metrics.Accuracy()
def dataset_fn(_):
return test_utils_obj.dataset_fn(feature_mapper, label_mapper)
@def_function.function
def train_step(iterator):
"""The step function for one training step."""
def step_fn(inputs):
"""The computation to run on each worker."""
features, labels = inputs
with backprop.GradientTape() as tape:
pred = model(features, training=True)
loss = keras.losses.binary_crossentropy(labels, pred)
loss = nn.compute_average_loss(loss)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(
list(zip(grads, model.trainable_variables)))
actual_pred = math_ops.cast(math_ops.greater(pred, 0.5),
dtypes.int64)
accuracy.update_state(labels, actual_pred)
strategy.run(step_fn, args=(next(iterator), ))
distributed_dataset = strategy.distribute_datasets_from_function(
dataset_fn)
distributed_iterator = iter(distributed_dataset)
num_epochs = 4
num_steps = 7
for _ in range(num_epochs):
accuracy.reset_state()
for _ in range(num_steps):
train_step(distributed_iterator)
self.assertGreater(accuracy.result().numpy(), 0.5)
self.assertEqual(optimizer.iterations.numpy(),
num_epochs * num_steps)
# Test save/load/serving the trained model.
test_utils_obj.test_save_load_serving_model(
model, feature_mapper,
test_utils_obj.define_reverse_lookup_layer())
class KerasMultiWorkerTestIndependentWorker(test.TestCase,
parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(
mode=['eager'],
strategy=[
strategy_combinations.multi_worker_mirrored_2x1_cpu,
strategy_combinations.multi_worker_mirrored_2x1_gpu,
]))
def testSimpleModelIndependentWorkerSync(self, strategy):
verification_callback = MultiWorkerVerificationCallback(
num_epoch=2,
num_worker=len(
json.loads(os.environ['TF_CONFIG'])['cluster']['worker']))
verification_callback.is_between_graph = \
strategy.extended.experimental_between_graph
batch_size = 64
steps = 2
train_ds, _ = multi_worker_testing_utils.mnist_synthetic_dataset(
batch_size, steps)
with strategy.scope():
model = multi_worker_testing_utils.get_mnist_model((28, 28, 1))
orig_loss, _ = model.evaluate(train_ds, steps=steps)
history = model.fit(x=train_ds,
epochs=2,
steps_per_epoch=steps,
callbacks=[verification_callback])
self.assertIsInstance(history, keras.callbacks.History)
trained_loss, _ = model.evaluate(train_ds, steps=steps)
self.assertLess(trained_loss, orig_loss)
verification_callback.verify(self)
class TestDistributionStrategyWithLossMasking(test.TestCase,
parameterized.TestCase):
# TODO(priyag): Enable all strategies for this test. Currently it does not
# work for TPU due to some invalid datatype.
@ds_combinations.generate(
combinations.combine(
distribution=[
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
],
mode=['graph', 'eager'],
optimizer=optimizer_combinations
.gradient_descent_optimizer_keras_v2_fn
))
def test_masking(self, distribution, optimizer):
with self.cached_session():
np.random.seed(1337)
x = np.array([[[1], [1]], [[0], [0]]])
with distribution.scope():
model = keras.models.Sequential()
model.add(keras.layers.Masking(mask_value=0, input_shape=(2, 1)))
model.add(
keras.layers.TimeDistributed(
keras.layers.Dense(1, kernel_initializer='one')))
model.compile(
loss='mse',
optimizer=optimizer())
y = np.array([[[1], [1]], [[1], [1]]])
dataset = dataset_ops.Dataset.from_tensor_slices((x, y))
dataset = dataset.repeat(100)
dataset = dataset.batch(10)
hist = model.fit(x=dataset, epochs=1, steps_per_epoch=2)
self.assertEqual(hist.history['loss'][0], 0)
class ModelCheckpointTest(test_base.IndependentWorkerTestBase,
parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(
mode=['graph'],
required_gpus=[0, 1],
file_format=['h5', 'tf'],
save_weights_only=[True, False]))
def testCheckpointExists(self, file_format, save_weights_only):
with self.cached_session():
train_ds, _ = multi_worker_testing_utils.mnist_synthetic_dataset(64, 2)
model = multi_worker_testing_utils.get_mnist_model((28, 28, 1))
saving_dir = self.get_temp_dir()
saving_filepath = os.path.join(saving_dir, 'checkpoint.' + file_format)
callbacks_list = [
callbacks.ModelCheckpoint(
filepath=saving_filepath, save_weights_only=save_weights_only)
]
self.assertFalse(file_io.file_exists_v2(saving_filepath))
model.fit(
x=train_ds, epochs=2, steps_per_epoch=2, callbacks=callbacks_list)
tf_saved_model_exists = file_io.file_exists_v2(saving_filepath)
tf_weights_only_checkpoint_exists = file_io.file_exists_v2(
saving_filepath + '.index')
self.assertTrue(
tf_saved_model_exists or tf_weights_only_checkpoint_exists)
class KerasModelsTest(test.TestCase, parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(distribution=strategy_combinations.all_strategies,
mode=["eager"]))
def test_lstm_model_with_dynamic_batch(self, distribution):
input_data = np.random.random([1, 32, 64, 64, 3])
input_shape = tuple(input_data.shape[1:])
def build_model():
model = keras.models.Sequential()
model.add(
keras.layers.ConvLSTM2D(4,
kernel_size=(4, 4),
activation="sigmoid",
padding="same",
input_shape=input_shape))
model.add(keras.layers.GlobalMaxPooling2D())
model.add(keras.layers.Dense(2, activation="sigmoid"))
return model
with distribution.scope():
model = build_model()
model.compile(loss="binary_crossentropy", optimizer="adam")
result = model.predict(input_data)
self.assertEqual(result.shape, (1, 2))
class LocalCollectiveAllReduceStrategy(
CollectiveAllReduceStrategyTestBase,
strategy_test_lib.DistributionTestBase,
strategy_test_lib.TwoDeviceDistributionTestBase,
parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(mode=['graph'], required_gpus=[2, 4]))
def testComplexModel(self, required_gpus):
self._test_complex_model(None, None, required_gpus)
@ds_combinations.generate(
combinations.combine(mode=['graph'], required_gpus=[2, 4]))
@testing_utils.enable_v2_dtype_behavior
def testMixedPrecision(self, required_gpus):
with policy.policy_scope('mixed_float16'):
self._test_mixed_precision(None, None, required_gpus)
def all_combinations():
return combinations.combine(distribution=[
strategy_combinations.default_strategy,
strategy_combinations.one_device_strategy,
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
strategy_combinations.mirrored_strategy_with_two_gpus,
],
mode=["graph"])
def test_combine_single_parameter(self):
self.assertEqual([{
"a": 1,
"b": 2
}, {
"a": 2,
"b": 2
}], combinations.combine(a=[1, 2], b=2))
def distributions_and_v1_and_v2_optimizers():
"""A common set of combination with DistributionStrategies and Optimizers."""
return combinations.combine(distribution=[
strategy_combinations_base.one_device_strategy,
strategy_combinations_base.mirrored_strategy_with_gpu_and_cpu,
strategy_combinations_base.mirrored_strategy_with_two_gpus,
],
optimizer_fn=optimizers_v1_and_v2)
def test_combinations_with_tpu_strategies_graph():
tpu_strategies = [
strategy_combinations.tpu_strategy,
]
return (combinations.times(
combinations.combine(distribution=tpu_strategies),
graph_mode_test_configuration()))
def strategy_combinations_eager_data_fn():
return combinations.combine(distribution=[
strategy_combinations.default_strategy,
strategy_combinations.one_device_strategy,
strategy_combinations.one_device_strategy_gpu,
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
strategy_combinations.mirrored_strategy_with_two_gpus
],
mode=['eager'],
data_fn=[get_numpy, get_dataset])
def test_times(self):
c1 = combinations.combine(mode=["graph"], loss=["callable", "tensor"])
c2 = combinations.combine(mode=["eager"], loss=["callable"])
c3 = combinations.combine(distribution=["d1", "d2"])
c4 = combinations.times(c3, c1 + c2)
self.assertEqual([
OrderedDict([("distribution", "d1"), ("loss", "callable"),
("mode", "graph")]),
OrderedDict([("distribution", "d1"), ("loss", "tensor"),
("mode", "graph")]),
OrderedDict([("distribution", "d1"), ("loss", "callable"),
("mode", "eager")]),
OrderedDict([("distribution", "d2"), ("loss", "callable"),
("mode", "graph")]),
OrderedDict([("distribution", "d2"), ("loss", "tensor"),
("mode", "graph")]),
OrderedDict([("distribution", "d2"), ("loss", "callable"),
("mode", "eager")])
], c4)
class KerasSaveLoadTest(test_base.TestSavedModelBase):
def setUp(self):
self._root_dir = 'keras_save_load'
super(KerasSaveLoadTest, self).setUp()
def _save_model(self, model, saved_dir):
model.save(saved_dir, save_format='tf')
def _load_and_run_model(self,
distribution,
saved_dir,
predict_dataset,
output_name='output_1'):
restored_keras_model = save.load_model(saved_dir)
return restored_keras_model.predict(predict_dataset,
steps=test_base.PREDICT_STEPS)
@ds_combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution):
self.run_test_save_no_strategy_restore_strategy(
model_and_input, distribution)
@ds_combinations.generate(
combinations.times(test_base.simple_models_with_strategies(),
combinations.combine(save_in_scope=[True, False])))
def test_save_strategy_restore_no_strategy(self, model_and_input,
distribution, save_in_scope):
self.run_test_save_strategy_restore_no_strategy(
model_and_input, distribution, save_in_scope)
@ds_combinations.generate(
combinations.times(test_base.simple_models_with_strategy_pairs(),
combinations.combine(save_in_scope=[True, False])))
def test_save_strategy_restore_strategy(self, model_and_input,
distribution_for_saving,
distribution_for_restoring,
save_in_scope):
self.run_test_save_strategy_restore_strategy(
model_and_input, distribution_for_saving,
distribution_for_restoring, save_in_scope)
class SingleLossStepTest(test.TestCase, parameterized.TestCase):
@ds_combinations.generate(
combinations.times(
optimizer_combinations.distributions_and_v1_optimizers(),
combinations.combine(
mode=strategy_combinations.graph_and_eager_modes),
combinations.combine(is_tpu=[False])) +
combinations.combine(distribution=[strategy_combinations.tpu_strategy],
optimizer_fn=optimizer_combinations.optimizers_v1,
mode=["graph"],
is_tpu=[True]))
def testTrainNetwork(self, distribution, optimizer_fn, is_tpu):
with distribution.scope():
single_loss_step, layer = single_loss_example(
optimizer_fn,
distribution,
use_bias=True,
iterations_per_step=2)
if context.executing_eagerly():
single_loss_step.initialize()
run_step = single_loss_step
else:
with self.cached_session() as sess:
sess.run(single_loss_step.initialize())
run_step = sess.make_callable(single_loss_step())
self.evaluate(variables.global_variables_initializer())
weights, biases = [], []
for _ in range(5):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
error = abs(
numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
class TestDistributionStrategyWithStaticShapes(test.TestCase,
parameterized.TestCase):
@ds_combinations.generate(
combinations.combine(
distribution=[
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
],
mode=['graph', 'eager']))
def test_input_batch_size_not_divisible_by_num_replicas(self, distribution):
with distribution.scope():
with self.assertRaisesRegex(
ValueError, r'The `batch_size` argument \(5\) must be divisible by '
r'the number of replicas \(2\)'):
keras.layers.Input(shape=(3,), batch_size=5, name='input')
@ds_combinations.generate(
combinations.combine(
distribution=[
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
],
mode=['graph', 'eager']))
def test_static_input_batch_size(self, distribution):
inputs = np.zeros((10, 3), dtype=np.float32)
targets = np.zeros((10, 4), dtype=np.float32)
dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
dataset = dataset.repeat(100)
dataset = dataset.batch(10, drop_remainder=True)
with distribution.scope():
x = keras.layers.Input(shape=(3,), batch_size=10, name='input')
y = keras.layers.Dense(4, name='dense')(x)
model = keras.Model(x, y)
model.compile(optimizer='sgd', loss='mse', metrics=['mae'])
model.fit(dataset, epochs=1, steps_per_epoch=5)
model.evaluate(dataset, steps=5)
model.predict(dataset)
def test_all_tf_execution_regimes(test_class_or_method=None):
"""Decorator for generating a collection of tests in various contexts.
Must be applied to subclasses of `parameterized.TestCase` (from
`absl/testing`), or a method of such a subclass.
When applied to a test method, this decorator results in the replacement of
that method with a collection of new test methods, each executed under a
different set of context managers that control some aspect of the execution
model. This decorator generates three test scenario combinations:
1. Eager mode with `tf.function` decorations enabled
2. Eager mode with `tf.function` decorations disabled
3. Graph mode (eveything)
When applied to a test class, all the methods in the class are affected.
Args:
test_class_or_method: the `TestCase` class or method to decorate.
Returns:
decorator: A generated TF `test_combinations` decorator, or if
`test_class_or_method` is not `None`, the generated decorator applied to
that function.
"""
decorator = test_combinations.generate(
(test_combinations.combine(mode='graph',
tf_function='enabled') +
test_combinations.combine(mode='eager',
tf_function=['enabled', 'disabled'])),
test_combinations=[
combinations.EagerGraphCombination(),
ExecuteFunctionsEagerlyCombination(),
])
if test_class_or_method:
return decorator(test_class_or_method)
return decorator
def test_combine(self):
self.assertEqual([{
"a": 1,
"b": 2
}, {
"a": 1,
"b": 3
}, {
"a": 2,
"b": 2
}, {
"a": 2,
"b": 3
}], combinations.combine(a=[1, 2], b=[2, 3]))
def strategy_combinations_eager_data_fn():
return combinations.combine(
distribution=[
strategy_combinations.default_strategy,
strategy_combinations.one_device_strategy,
strategy_combinations.one_device_strategy_gpu,
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
strategy_combinations.mirrored_strategy_with_two_gpus,
strategy_combinations.multi_worker_mirrored_2x1_cpu,
strategy_combinations.multi_worker_mirrored_2x1_gpu,
strategy_combinations.multi_worker_mirrored_2x2_gpu,
# NOTE: TPUStrategy not tested because the models in this test are
# sparse and do not work with TPUs.
],
mode=['eager'],
data_fn=['numpy', 'dataset'])
