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)