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)
@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)
@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)
@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)
@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)
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, experimental_run_tf_function):
return test_base.load_and_run_with_saved_model_api(distribution, saved_dir,
predict_dataset,
output_name)
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution,
experimental_run_tf_function):
self.run_test_save_no_strategy_restore_strategy(
model_and_input, distribution, experimental_run_tf_function)
@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,
experimental_run_tf_function):
if save_in_scope:
# TODO(b/134703272): Unskip this test when fixed.
self.skipTest(('Saving model within tf.distribute.Strategy scope is not ',
'supported.'))
self.run_test_save_strategy_restore_no_strategy(
model_and_input, distribution, save_in_scope,
experimental_run_tf_function)
@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,
experimental_run_tf_function):
if save_in_scope:
# TODO(b/134703272): Unskip this test when fixed.
self.skipTest(('Saving model within tf.distribute.Strategy scope is not ',
'supported.'))
self.run_test_save_strategy_restore_strategy(model_and_input,
distribution_for_saving,
distribution_for_restoring,
save_in_scope,
experimental_run_tf_function)
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, experimental_run_tf_function):
restored_keras_model = save.load_model(saved_dir)
restored_keras_model._experimental_run_tf_function = (
experimental_run_tf_function)
return restored_keras_model.predict(
predict_dataset, steps=test_base.PREDICT_STEPS)
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution,
experimental_run_tf_function):
self.run_test_save_no_strategy_restore_strategy(
model_and_input, distribution, experimental_run_tf_function)
@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,
experimental_run_tf_function):
if save_in_scope:
self.skipTest(('b/134703272 - Saving model in tf.distribute.Strategy ',
'scope is not supported.'))
self.run_test_save_strategy_restore_no_strategy(
model_and_input, distribution, save_in_scope,
experimental_run_tf_function)
@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,
experimental_run_tf_function):
if save_in_scope:
self.skipTest(('b/134703272 - Saving model in tf.distribute.Strategy ',
'scope is not supported.'))
self.run_test_save_strategy_restore_strategy(model_and_input,
distribution_for_saving,
distribution_for_restoring,
save_in_scope,
experimental_run_tf_function)
class SavedModelSaveAndLoadTest(test_base.TestSavedModelBase):
def setUp(self):
self._root_dir = 'saved_model_save_load'
super(SavedModelSaveAndLoadTest, 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):
func = saved_model.load(saved_dir)
if distribution:
dist_predict_dataset = distribution.experimental_distribute_dataset(
predict_dataset)
per_replica_predict_data = next(iter(dist_predict_dataset))
result = distribution.experimental_run_v2(
func.signatures[_DEFAULT_FUNCTION_KEY],
args=(per_replica_predict_data, ))
result = result[output_name]
# Convert the per_replica value to a list, then concatenate them
reduced = distribution.experimental_local_results(result)
concat = array_ops.concat(reduced, 0)
return concat
else:
result = func.signatures[_DEFAULT_FUNCTION_KEY](next(
iter(predict_dataset)))
return result[output_name]
@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)
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_strategy_restore_no_strategy(self, model_and_input,
distribution):
self.skipTest(
('Saving/loading model with tf.distribute.Strategy is not ',
'supported.'))
self.run_test_save_strategy_restore_no_strategy(
model_and_input, distribution)
@combinations.generate(test_base.simple_models_with_strategy_pairs())
def test_save_strategy_restore_strategy(self, model_and_input,
distribution_pair):
self.skipTest(
('Saving/loading model with tf.distribute.Strategy is not ',
'supported.'))
self.run_test_save_strategy_restore_strategy(model_and_input,
distribution_pair)
class SavedModelSaveAndLoadTest(test_base.TestSavedModelBase):
def setUp(self):
self._root_dir = 'saved_model_save_load'
super(SavedModelSaveAndLoadTest, self).setUp()
def _save_model(self, model, saved_dir):
keras_saved_model.export_saved_model(model, saved_dir, serving_only=True)
def _load_and_run_model(self, distribution, saved_dir, predict_dataset,
output_name, run_distributed):
return test_base.load_and_run_with_saved_model_api(distribution, saved_dir,
predict_dataset,
output_name)
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution, run_distributed):
self.run_test_save_no_strategy_restore_strategy(model_and_input,
distribution,
run_distributed)
@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,
run_distributed):
if save_in_scope:
self.skipTest(('Saving model within tf.distribute.Strategy scope is not ',
'supported.'))
self.run_test_save_strategy_restore_no_strategy(model_and_input,
distribution, save_in_scope,
run_distributed)
@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, run_distributed):
if save_in_scope:
self.skipTest(('Saving model within tf.distribute.Strategy scope is not ',
'supported.'))
self.run_test_save_strategy_restore_strategy(model_and_input,
distribution_for_saving,
distribution_for_restoring,
save_in_scope, run_distributed)
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,
experimental_run_tf_function,
output_name='output_1'):
return test_base.load_and_run_with_saved_model_api(distribution, saved_dir,
predict_dataset,
output_name)
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution,
experimental_run_tf_function):
self.run_test_save_no_strategy_restore_strategy(
model_and_input, distribution, experimental_run_tf_function)
@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,
experimental_run_tf_function):
self.run_test_save_strategy_restore_no_strategy(
model_and_input, distribution, save_in_scope,
experimental_run_tf_function)
@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,
experimental_run_tf_function):
self.run_test_save_strategy_restore_strategy(model_and_input,
distribution_for_saving,
distribution_for_restoring,
save_in_scope,
experimental_run_tf_function)
class KerasExperimentalSaveLoadTest(test_base.TestSavedModelBase):
def setUp(self):
self._root_dir = 'keras_experimental_save_load'
super(KerasExperimentalSaveLoadTest, self).setUp()
def _save_model(self, model, saved_dir):
saved_model.export_saved_model(model, saved_dir)
def _load_and_run_model(self, distribution, saved_dir, predict_dataset,
output_name, experimental_run_tf_function):
restored_keras_model = saved_model.load_from_saved_model(saved_dir)
restored_keras_model._experimental_run_tf_function = (
experimental_run_tf_function)
return restored_keras_model.predict(predict_dataset,
steps=test_base.PREDICT_STEPS)
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution,
experimental_run_tf_function):
self.run_test_save_no_strategy_restore_strategy(
model_and_input, distribution, experimental_run_tf_function)
@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,
experimental_run_tf_function):
self.run_test_save_strategy_restore_no_strategy(
model_and_input, distribution, save_in_scope,
experimental_run_tf_function)
@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,
experimental_run_tf_function):
self.run_test_save_strategy_restore_strategy(
model_and_input, distribution_for_saving,
distribution_for_restoring, save_in_scope,
experimental_run_tf_function)
class SavedModelSaveAndLoadTest(test_base.TestSavedModelBase):
def setUp(self):
self._root_dir = 'saved_model_save_load'
super(SavedModelSaveAndLoadTest, 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):
dist_predict_dataset = distribution.experimental_distribute_dataset(
predict_dataset)
per_replica_predict_data = next(iter(dist_predict_dataset))
func = saved_model.load(saved_dir)
result = distribution.experimental_run_v2(
func.signatures[_DEFAULT_FUNCTION_KEY], per_replica_predict_data)
return result[output_name]
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_no_strategy_restore_strategy(self, model_and_input,
distribution):
self.skipTest(('Saving/loading model with tf.distribute.Strategy is not ',
'supported.'))
self.run_test_save_no_strategy_restore_strategy(model_and_input,
distribution)
@combinations.generate(test_base.simple_models_with_strategies())
def test_save_strategy_restore_no_strategy(self, model_and_input,
distribution):
self.skipTest(('Saving/loading model with tf.distribute.Strategy is not ',
'supported.'))
self.run_test_save_strategy_restore_no_strategy(model_and_input,
distribution)
@combinations.generate(test_base.simple_models_with_strategy_pairs())
def test_save_strategy_restore_strategy(self, model_and_input,
distribution_pair):
self.skipTest(('Saving/loading model with tf.distribute.Strategy is not ',
'supported.'))
self.run_test_save_strategy_restore_strategy(model_and_input,
distribution_pair)
