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): tf.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) @tf.__internal__.distribute.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) @tf.__internal__.distribute.combinations.generate( tf.__internal__.test.combinations.times(test_base.simple_models_with_strategies(), tf.__internal__.test.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) @tf.__internal__.distribute.combinations.generate( tf.__internal__.test.combinations.times(test_base.simple_models_with_strategy_pairs(), tf.__internal__.test.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) @tf.__internal__.distribute.combinations.generate( tf.__internal__.test.combinations.times(test_base.simple_models_with_strategies(), tf.__internal__.test.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 = tf.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 KerasSaveLoadTest(test_base.TestSavedModelBase): def setUp(self): self._root_dir = "keras_save_load" super().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) @tf.__internal__.distribute.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) @tf.__internal__.distribute.combinations.generate( tf.__internal__.test.combinations.times( test_base.simple_models_with_strategies(), tf.__internal__.test.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) @tf.__internal__.distribute.combinations.generate( tf.__internal__.test.combinations.times( test_base.simple_models_with_strategy_pairs(), tf.__internal__.test.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 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): save.save_model(model, saved_dir, save_format='tf') 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) @tf.__internal__.distribute.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) @tf.__internal__.distribute.combinations.generate( tf.__internal__.test.combinations.times( test_base.simple_models_with_strategies(), tf.__internal__.test.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) @tf.__internal__.distribute.combinations.generate( tf.__internal__.test.combinations.times( test_base.simple_models_with_strategy_pairs(), tf.__internal__.test.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)