Fix save model related unit test.

1. The save_test requires cached_session() in graph mode.
2. Callback test is failing in OSS the model output shape doesn't match with the input data. Not sure why its not failing in python/keras.

PiperOrigin-RevId: 360262858
Change-Id: Ie5943a83943d1d759f061c41bf8877c8fd6a0506
This commit is contained in:
Scott Zhu 2021-03-01 13:23:21 -08:00 committed by TensorFlower Gardener
parent b934c56a68
commit 21745f6a97
2 changed files with 21 additions and 20 deletions

View File

@ -751,6 +751,7 @@ class KerasCallbacksTest(keras_parameterized.TestCase):
test_samples=TEST_SAMPLES,
input_shape=(INPUT_DIM,),
num_classes=NUM_CLASSES)
y_train = np_utils.to_categorical(y_train, num_classes=NUM_CLASSES)
model.fit(
x_train,
@ -760,7 +761,6 @@ class KerasCallbacksTest(keras_parameterized.TestCase):
verbose=0)
# Check that the filepath is a SavedModel directory.
self.assertIn('saved_model.pb', os.listdir(filepath))
os.remove(filepath)
def _get_dummy_resource_for_model_checkpoint_testing(self):

View File

@ -1012,27 +1012,28 @@ class TestWholeModelSaving(keras_parameterized.TestCase):
e.g. "head_0_accuracy" should not become "head_0_head_0_accuracy" after
saving and loading a model.
"""
input_ = keras.Input((4,))
model = keras.Model(
input_,
[keras.layers.Softmax(name='head_0')(keras.layers.Dense(3)(input_)),
keras.layers.Softmax(name='head_1')(keras.layers.Dense(5)(input_))])
metric = keras.metrics.BinaryAccuracy()
model.compile(optimizer='rmsprop',
loss='mse',
metrics={'head_0': [metric, 'accuracy']})
with self.cached_session():
input_ = keras.Input((4,))
model = keras.Model(
input_,
[keras.layers.Softmax(name='head_0')(keras.layers.Dense(3)(input_)),
keras.layers.Softmax(name='head_1')(keras.layers.Dense(5)(input_))])
metric = keras.metrics.BinaryAccuracy()
model.compile(optimizer='rmsprop',
loss='mse',
metrics={'head_0': [metric, 'accuracy']})
# Run one iteration.
x = np.random.rand(2, 4)
y = {'head_0': np.random.randint(2, size=(2, 3)),
'head_1': np.random.randint(2, size=(2, 5))}
model.fit(x, y, verbose=0)
# Run one iteration.
x = np.random.rand(2, 4)
y = {'head_0': np.random.randint(2, size=(2, 3)),
'head_1': np.random.randint(2, size=(2, 5))}
model.fit(x, y, verbose=0)
# Save and reload.
save_format = testing_utils.get_save_format()
saved_model_dir = self._save_model_dir()
keras.models.save_model(model, saved_model_dir, save_format=save_format)
loaded = keras.models.load_model(saved_model_dir)
# Save and reload.
save_format = testing_utils.get_save_format()
saved_model_dir = self._save_model_dir()
keras.models.save_model(model, saved_model_dir, save_format=save_format)
loaded = keras.models.load_model(saved_model_dir)
# Make sure the metrics names from the model before saving match the loaded
# model.