experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Migrate saved_model_mixed_api_test to V2 API

Browse Source 
PiperOrigin-RevId: 315930643
Change-Id: I388b88c863fd01b2077c523de3c315487b4f2a04
This commit is contained in:
Ran Chen 2020-06-11 10:43:21 -07:00 committed by TensorFlower Gardener
parent d394de5e90
commit 863ceb97de
2 changed files with 4 additions and 3 deletions
tensorflow/python/distribute
BUILDsaved_model_mixed_api_test.py

3
tensorflow/python/distribute/BUILD
View File

@ -1597,9 +1597,10 @@ distribute_py_test(
"no_rocm",
],
deps = [
":combinations",
":saved_model_test_base",
"//tensorflow/python/eager:test",
"//tensorflow/python/keras/saving",
"//tensorflow/python/saved_model",
],
)

4
tensorflow/python/distribute/saved_model_mixed_api_test.py
View File

@ -27,7 +27,7 @@ from __future__ import print_function
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import saved_model_test_base as test_base
from tensorflow.python.eager import test
from tensorflow.python.keras.saving import saved_model_experimental as keras_saved_model
from tensorflow.python.keras.saving import save
_DEFAULT_FUNCTION_KEY = 'serving_default'
@ -39,7 +39,7 @@ class SavedModelSaveAndLoadTest(test_base.TestSavedModelBase):
super(SavedModelSaveAndLoadTest, self).setUp()
def _save_model(self, model, saved_dir):
keras_saved_model.export_saved_model(model, saved_dir, serving_only=True)
save.save_model(model, saved_dir, save_format='tf')
def _load_and_run_model(self,
distribution,