def test_vocabulary_persistence_across_saving(self):
vocab_data = ["earth", "wind", "and", "fire"]
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
# Build and validate a golden model.
input_data = keras.Input(shape=(None, ), dtype=dtypes.string)
layer = get_layer_class()(max_tokens=None,
standardize=None,
split=None,
output_mode=text_vectorization.INT)
layer.set_vocabulary(vocab_data)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_array)
self.assertAllEqual(output_dataset, expected_output)
# Save the model to disk.
output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model")
saving.export_saved_model(model, output_path)
loaded_model = saving.load_from_saved_model(
output_path,
custom_objects={"TextVectorization": get_layer_class()})
# Ensure that the loaded model is unique (so that the save/load is real)
self.assertIsNot(model, loaded_model)
# Validate correctness of the new model.
new_output_dataset = loaded_model.predict(input_array)
self.assertAllEqual(new_output_dataset, expected_output)
def test_saving_sequential_model(self):
with self.cached_session():
model = keras.models.Sequential()
model.add(keras.layers.Dense(2, input_shape=(3, )))
model.add(keras.layers.RepeatVector(3))
model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
model.compile(loss=keras.losses.MSE,
optimizer=rmsprop.RMSprop(lr=0.0001),
metrics=[keras.metrics.categorical_accuracy],
sample_weight_mode='temporal',
run_eagerly=testing_utils.should_run_eagerly(),
experimental_run_tf_function=testing_utils.
should_run_tf_function())
x = np.random.random((1, 3))
y = np.random.random((1, 3, 3))
model.train_on_batch(x, y)
ref_y = model.predict(x)
saved_model_dir = self._save_model_dir()
keras_saved_model.export_saved_model(model, saved_model_dir)
loaded_model = keras_saved_model.load_from_saved_model(
saved_model_dir)
y = loaded_model.predict(x)
self.assertAllClose(ref_y, y, atol=1e-05)
def test_saving_functional_model(self):
with self.cached_session():
inputs = keras.layers.Input(shape=(3, ))
x = keras.layers.Dense(2)(inputs)
output = keras.layers.Dense(3)(x)
model = keras.models.Model(inputs, output)
model.compile(loss=keras.losses.MSE,
optimizer=rmsprop.RMSprop(lr=0.0001),
metrics=[keras.metrics.categorical_accuracy],
run_eagerly=testing_utils.should_run_eagerly(),
experimental_run_tf_function=testing_utils.
should_run_tf_function())
x = np.random.random((1, 3))
y = np.random.random((1, 3))
model.train_on_batch(x, y)
ref_y = model.predict(x)
saved_model_dir = self._save_model_dir()
keras_saved_model.export_saved_model(model, saved_model_dir)
loaded_model = keras_saved_model.load_from_saved_model(
saved_model_dir)
y = loaded_model.predict(x)
self.assertAllClose(ref_y, y, atol=1e-05)
def _save_restore_saved_model(model):
tmpdir = tempfile.mkdtemp()
saved_model_experimental.export_saved_model(model, tmpdir)
with prune.prune_scope():
loaded_model = saved_model_experimental.load_from_saved_model(tmpdir)
loaded_model.compile(loss='categorical_crossentropy',
optimizer='sgd',
metrics=['accuracy'])
return loaded_model
def test_saving_sequential_model_without_compile(self):
with self.cached_session():
model = keras.models.Sequential()
model.add(keras.layers.Dense(2, input_shape=(3,)))
model.add(keras.layers.RepeatVector(3))
model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
x = np.random.random((1, 3))
ref_y = model.predict(x)
saved_model_dir = self._save_model_dir()
keras_saved_model.export_saved_model(model, saved_model_dir)
loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)
y = loaded_model.predict(x)
self.assertAllClose(ref_y, y, atol=1e-05)
def testSaveSequentialModelWithoutInputShapes(self):
model = sequential_model_without_input_shape(True)
# A Sequential model that hasn't been built should raise an error.
with self.assertRaisesRegex(
ValueError, 'Weights for sequential model have not yet been created'):
keras_saved_model.export_saved_model(model, '')
# Even with input_signature, the model's weights has not been created.
with self.assertRaisesRegex(
ValueError, 'Weights for sequential model have not yet been created'):
saved_model_dir = self._save_model_dir()
keras_saved_model.export_saved_model(
model,
saved_model_dir,
input_signature=tensor_spec.TensorSpec(
shape=(10, 11, 12, 13, 14), dtype=dtypes.float32,
name='spec_input'))
def test_saving_subclassed_model_raise_error(self):
# For now, saving subclassed model should raise an error. It should be
# avoided later with loading from SavedModel.pb.
class SubclassedModel(model_lib.Model):
def __init__(self):
super(SubclassedModel, self).__init__()
self.layer1 = keras.layers.Dense(3)
self.layer2 = keras.layers.Dense(1)
def call(self, inp):
return self.layer2(self.layer1(inp))
model = SubclassedModel()
saved_model_dir = self._save_model_dir()
with self.assertRaises(NotImplementedError):
keras_saved_model.export_saved_model(model, saved_model_dir)
def testSaveAndLoadSavedModelWithCustomObject(self):
saved_model_dir = self._save_model_dir()
with session.Session(graph=ops.Graph()) as sess:
def relu6(x):
return keras.backend.relu(x, max_value=6)
inputs = keras.layers.Input(shape=(1,))
outputs = keras.layers.Activation(relu6)(inputs)
model = keras.models.Model(inputs, outputs)
keras_saved_model.export_saved_model(
model, saved_model_dir, custom_objects={'relu6': relu6})
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, _ = load_model(sess, saved_model_dir,
mode_keys.ModeKeys.PREDICT)
input_name = model.input_names[0]
output_name = model.output_names[0]
predictions = sess.run(
outputs[output_name], {inputs[input_name]: [[7], [-3], [4]]})
self.assertAllEqual([[6], [0], [4]], predictions)
def test_saving_with_tf_optimizer(self):
model = keras.models.Sequential()
model.add(keras.layers.Dense(2, input_shape=(3, )))
model.add(keras.layers.Dense(3))
model.compile(loss='mse',
optimizer=training_module.RMSPropOptimizer(0.1),
metrics=['acc'])
x = np.random.random((1, 3))
y = np.random.random((1, 3))
model.train_on_batch(x, y)
ref_y = model.predict(x)
saved_model_dir = self._save_model_dir()
keras_saved_model.export_saved_model(model, saved_model_dir)
loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)
loaded_model.compile(loss='mse',
optimizer=training_module.RMSPropOptimizer(0.1),
metrics=['acc'],
run_eagerly=testing_utils.should_run_eagerly(),
experimental_run_tf_function=testing_utils.
should_run_tf_function())
y = loaded_model.predict(x)
self.assertAllClose(ref_y, y, atol=1e-05)
# test that new updates are the same with both models
x = np.random.random((1, 3))
y = np.random.random((1, 3))
ref_loss = model.train_on_batch(x, y)
loss = loaded_model.train_on_batch(x, y)
self.assertAllClose(ref_loss, loss, atol=1e-05)
ref_y = model.predict(x)
y = loaded_model.predict(x)
self.assertAllClose(ref_y, y, atol=1e-05)
# test saving/loading again
saved_model_dir2 = self._save_model_dir('saved_model_2')
keras_saved_model.export_saved_model(loaded_model, saved_model_dir2)
loaded_model = keras_saved_model.load_from_saved_model(
saved_model_dir2)
y = loaded_model.predict(x)
self.assertAllClose(ref_y, y, atol=1e-05)
def test_saving_functional_model_without_compile(self):
with self.cached_session():
inputs = keras.layers.Input(shape=(3,))
x = keras.layers.Dense(2)(inputs)
output = keras.layers.Dense(3)(x)
model = keras.models.Model(inputs, output)
x = np.random.random((1, 3))
y = np.random.random((1, 3))
ref_y = model.predict(x)
saved_model_dir = self._save_model_dir()
keras_saved_model.export_saved_model(model, saved_model_dir)
loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)
y = loaded_model.predict(x)
self.assertAllClose(ref_y, y, atol=1e-05)
def testServingOnly(self, model_builder, input_signature):
if context.executing_eagerly():
saved_model_dir = self._save_model_dir()
input_arr = np.random.random((5, 3)).astype(np.float32)
model = model_builder()
ref_predict = model.predict(input_arr)
keras_saved_model.export_saved_model(
model,
saved_model_dir,
serving_only=True,
input_signature=input_signature)
# Load predict graph, and test predictions
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, _ = load_model(sess, saved_model_dir,
mode_keys.ModeKeys.PREDICT)
predictions = sess.run(outputs[next(iter(outputs.keys()))],
{inputs[next(iter(inputs.keys()))]: input_arr})
self.assertAllClose(ref_predict, predictions, atol=1e-05)
def test_vocabulary_persistence_across_saving_with_tfidf(self):
vocab_data = ["earth", "wind", "and", "fire"]
tfidf_data = [.5, .25, .2, .125]
input_array = np.array([["earth", "wind", "and", "earth"],
["ohio", "fire", "earth", "michigan"]])
# pyformat: disable
# pylint: disable=bad-whitespace
expected_output = [[ 0, 1, .25, .2, 0],
[.1, .5, 0, 0, .125]]
# pylint: enable=bad-whitespace
# pyformat: enable
# Build and validate a golden model.
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = get_layer_class()(
max_tokens=5,
standardize=None,
split=None,
output_mode=text_vectorization.TFIDF)
layer.set_vocabulary(vocab_data, df_data=tfidf_data, oov_df_value=.05)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_array)
self.assertAllClose(output_dataset, expected_output)
# Save the model to disk.
output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model")
saving.export_saved_model(model, output_path)
loaded_model = saving.load_from_saved_model(
output_path, custom_objects={"TextVectorization": get_layer_class()})
# Ensure that the loaded model is unique (so that the save/load is real)
self.assertIsNot(model, loaded_model)
# Validate correctness of the new model.
new_output_dataset = loaded_model.predict(input_array)
self.assertAllClose(new_output_dataset, expected_output)
def _save_model(self, model, saved_dir):
saved_model.export_saved_model(model, saved_dir)
def testSaveAndLoadSavedModelExport(self, model_builder,
uses_learning_phase, optimizer_cls,
train_before_export):
optimizer = None if optimizer_cls is None else optimizer_cls()
saved_model_dir = self._save_model_dir()
np.random.seed(130)
input_arr = np.random.random((1, 3))
target_arr = np.random.random((1, 3))
model = model_builder(uses_learning_phase)
if optimizer is not None:
model.compile(loss='mse', optimizer=optimizer, metrics=['mae'])
if train_before_export:
model.train_on_batch(input_arr, target_arr)
ref_loss, ref_mae = model.evaluate(input_arr, target_arr)
ref_predict = model.predict(input_arr)
# Export SavedModel
keras_saved_model.export_saved_model(model, saved_model_dir)
input_name = model.input_names[0]
output_name = model.output_names[0]
target_name = output_name + '_target'
# Load predict graph, and test predictions
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, _ = load_model(sess, saved_model_dir,
mode_keys.ModeKeys.PREDICT)
predictions = sess.run(outputs[output_name],
{inputs[input_name]: input_arr})
self.assertAllClose(ref_predict, predictions, atol=1e-05)
if optimizer:
# Load eval graph, and test predictions, loss and metric values
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, _ = load_model(sess, saved_model_dir,
mode_keys.ModeKeys.TEST)
# First obtain the loss and predictions, and run the metric update op by
# feeding in the inputs and targets.
metrics_name = 'mae' if tf2.enabled(
) else 'mean_absolute_error'
metrics_update_op_key = 'metrics/' + metrics_name + '/update_op'
metrics_value_op_key = 'metrics/' + metrics_name + '/value'
loss, predictions, _ = sess.run(
(outputs['loss'], outputs['predictions/' + output_name],
outputs[metrics_update_op_key]), {
inputs[input_name]: input_arr,
inputs[target_name]: target_arr
})
# The metric value should be run after the update op, to ensure that it
# reflects the correct value.
metric_value = sess.run(outputs[metrics_value_op_key])
self.assertEqual(int(train_before_export),
sess.run(training_module.get_global_step()))
self.assertAllClose(ref_loss, loss, atol=1e-05)
self.assertAllClose(ref_mae, metric_value, atol=1e-05)
self.assertAllClose(ref_predict, predictions, atol=1e-05)
# Load train graph, and check for the train op, and prediction values
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, meta_graph_def = load_model(
sess, saved_model_dir, mode_keys.ModeKeys.TRAIN)
self.assertEqual(int(train_before_export),
sess.run(training_module.get_global_step()))
self.assertIn('loss', outputs)
self.assertIn(metrics_update_op_key, outputs)
self.assertIn(metrics_value_op_key, outputs)
self.assertIn('predictions/' + output_name, outputs)
# Train for a step
train_op = loader_impl.get_train_op(meta_graph_def)
train_outputs, _ = sess.run([outputs, train_op], {
inputs[input_name]: input_arr,
inputs[target_name]: target_arr
})
self.assertEqual(
int(train_before_export) + 1,
sess.run(training_module.get_global_step()))
if uses_learning_phase:
self.assertAllClose([[0, 0, 0]],
train_outputs['predictions/' +
output_name],
atol=1e-05)
else:
self.assertNotAllClose([[0, 0, 0]],
train_outputs['predictions/' +
output_name],
atol=1e-05)
def _save_model(self, model, saved_dir): keras_saved_model.export_saved_model(model, saved_dir, serving_only=True)
def save_model(model):
export_path = os.path.join(args.job_dir, 'keras_export', args.job_name)
export_saved_model(model, export_path)
# tf.keras.models.save_model(keras_model, export_path)
print('Model exported to: {}'.format(export_path))
