Export Keras model to SavedModel.

PiperOrigin-RevId: 210648154

4 files changed, 513 insertions, 64 deletions

diff --git a/tensorflow/contrib/saved_model/BUILD b/tensorflow/contrib/saved_model/BUILD
index e7eb4ac563..b897224c6d 100644
--- a/tensorflow/contrib/saved_model/BUILD
+++ b/tensorflow/contrib/saved_model/BUILD
@@ -36,6 +36,7 @@ py_library(
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
     deps = [
+        ":keras_saved_model",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:framework_ops",
         "//tensorflow/python:lib",
@@ -101,23 +102,33 @@ py_library(
     tags = ["no_windows"],
     visibility = ["//visibility:public"],
     deps = [
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:framework_ops",
         "//tensorflow/python:lib",
+        "//tensorflow/python:metrics",
+        "//tensorflow/python:platform",
+        "//tensorflow/python:saver",
         "//tensorflow/python:util",
+        "//tensorflow/python/estimator",
+        "//tensorflow/python/estimator:export",
+        "//tensorflow/python/estimator:keras",
+        "//tensorflow/python/estimator:model_fn",
         "//tensorflow/python/keras:engine",
-        "//tensorflow/python/saved_model:constants",
+        "//tensorflow/python/saved_model",
     ],
 )
 
 py_test(
     name = "keras_saved_model_test",
-    size = "small",
+    size = "medium",
     srcs = ["python/saved_model/keras_saved_model_test.py"],
     srcs_version = "PY2AND3",
     deps = [
-        ":saved_model_py",
+        ":keras_saved_model",
         "//tensorflow/python:client_testlib",
         "//tensorflow/python:training",
         "//tensorflow/python/keras",
         "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
     ],
 )
diff --git a/tensorflow/contrib/saved_model/__init__.py b/tensorflow/contrib/saved_model/__init__.py
index 95e1a8967b..074dc655ac 100644
--- a/tensorflow/contrib/saved_model/__init__.py
+++ b/tensorflow/contrib/saved_model/__init__.py
@@ -26,10 +26,13 @@ from __future__ import print_function
 # pylint: disable=unused-import,wildcard-import,line-too-long
 from tensorflow.contrib.saved_model.python.saved_model.keras_saved_model import *
 from tensorflow.contrib.saved_model.python.saved_model.signature_def_utils import *
-# pylint: enable=unused-import,widcard-import,line-too-long
+# pylint: enable=unused-import,wildcard-import,line-too-long
 
 from tensorflow.python.util.all_util import remove_undocumented
 
-_allowed_symbols = ["get_signature_def_by_key", "load_model", "save_model"]
+_allowed_symbols = [
+    "get_signature_def_by_key",
+    "load_keras_model",
+    "save_keras_model"]
 
 remove_undocumented(__name__, _allowed_symbols)
diff --git a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py
index e2a969f053..2c5c8c4afd 100644
--- a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py
+++ b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py
@@ -20,28 +20,69 @@ from __future__ import print_function
 
 import os
 
+from tensorflow.python.client import session
+from tensorflow.python.estimator import keras as estimator_keras_util
+from tensorflow.python.estimator import model_fn as model_fn_lib
+from tensorflow.python.estimator.export import export as export_helpers
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras import models as models_lib
+from tensorflow.python.keras import optimizers
 from tensorflow.python.keras.models import model_from_json
 from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import builder as saved_model_builder
 from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import utils_impl as saved_model_utils
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training.checkpointable import util as checkpointable_utils
 from tensorflow.python.util import compat
 
 
-def save_model(model, saved_model_path):
+def save_keras_model(
+    model, saved_model_path, custom_objects=None, as_text=None):
   """Save a `tf.keras.Model` into Tensorflow SavedModel format.
 
-  `save_model` generates such files/folders under the `saved_model_path` folder:
+  `save_model` generates new files/folders under the `saved_model_path` folder:
   1) an asset folder containing the json string of the model's
-  configuration(topology).
+     configuration (topology).
   2) a checkpoint containing the model weights.
+  3) a saved_model.pb file containing the model's MetaGraphs. The prediction
+     graph is always exported. The evaluaton and training graphs are exported
+     if the following conditions are met:
+     - Evaluation: model loss is defined.
+     - Training: model is compiled with an optimizer defined under `tf.train`.
+       This is because `tf.keras.optimizers.Optimizer` instances cannot be
+       saved to checkpoints.
 
-  Note that subclassed models can not be saved via this function, unless you
-  provide an implementation for get_config() and from_config().
-  Also note that `tf.keras.optimizers.Optimizer` instances can not currently be
-  saved to checkpoints. Use optimizers from `tf.train`.
+  Model Requirements:
+  - Model must be a sequential model or functional model. Subclassed models can
+    not be saved via this function, unless you provide an implementation for
+    get_config() and from_config().
+  - All variables must be saveable by the model. In general, this condition is
+    met through the use of layers defined in the keras library. However,
+    there is currently a bug with variables created in Lambda layer functions
+    not being saved correctly (see
+    https://github.com/keras-team/keras/issues/9740).
+
+  Note that each mode is exported in separate graphs, so different modes do not
+  share variables. To use the train graph with evaluation or prediction graphs,
+  create a new checkpoint if variable values have been updated.
 
   Args:
     model: A `tf.keras.Model` to be saved.
     saved_model_path: a string specifying the path to the SavedModel directory.
+      The SavedModel will be saved to a timestamped folder created within this
+      directory.
+    custom_objects: Optional dictionary mapping string names to custom classes
+      or functions (e.g. custom loss functions).
+    as_text: whether to write the `SavedModel` proto in text format.
+
+  Returns:
+    String path to the SavedModel folder, a subdirectory of `saved_model_path`.
 
   Raises:
     NotImplementedError: If the passed in model is a subclassed model.
@@ -49,35 +90,200 @@ def save_model(model, saved_model_path):
   if not model._is_graph_network:
     raise NotImplementedError
 
-  # save model configuration as a json string under assets folder.
-  model_json = model.to_json()
-  assets_destination_dir = os.path.join(
-      compat.as_bytes(saved_model_path),
-      compat.as_bytes(constants.ASSETS_DIRECTORY))
+  export_dir = export_helpers.get_timestamped_export_dir(saved_model_path)
+  temp_export_dir = export_helpers.get_temp_export_dir(export_dir)
+
+  builder = saved_model_builder.SavedModelBuilder(temp_export_dir)
+
+  # Manually save variables to export them in an object-based checkpoint. This
+  # skips the `builder.add_meta_graph_and_variables()` step, which saves a
+  # named-based checkpoint.
+  # TODO(b/113134168): Add fn to Builder to save with object-based saver.
+  # TODO(b/113178242): This should only export the model json structure. Only
+  # one save is needed once the weights can be copied from the model to clone.
+  checkpoint_path = _export_model_json_and_variables(model, temp_export_dir)
+
+  # Export each mode. Use ModeKeys enums defined for `Estimator` to ensure that
+  # Keras models and `Estimator`s are exported with the same format.
+  # Every time a mode is exported, the code checks to see if new variables have
+  # been created (e.g. optimizer slot variables). If that is the case, the
+  # checkpoint is re-saved to include the new variables.
+  export_args = {'builder': builder,
+                 'model': model,
+                 'custom_objects': custom_objects,
+                 'checkpoint_path': checkpoint_path}
+
+  has_saved_vars = False
+  if model.optimizer:
+    if isinstance(model.optimizer, optimizers.TFOptimizer):
+      _export_mode(model_fn_lib.ModeKeys.TRAIN, has_saved_vars, **export_args)
+      has_saved_vars = True
+      _export_mode(model_fn_lib.ModeKeys.EVAL, has_saved_vars, **export_args)
+    else:
+      logging.warning(
+          'Model was compiled with an optimizer, but the optimizer is not from '
+          '`tf.train` (e.g. `tf.train.AdagradOptimizer`). Only the serving '
+          'graph was exported. The train and evaluate graphs were not added to '
+          'the SavedModel.')
+  _export_mode(model_fn_lib.ModeKeys.PREDICT, has_saved_vars, **export_args)
+
+  builder.save(as_text)
+
+  gfile.Rename(temp_export_dir, export_dir)
+  return export_dir
 
-  if not file_io.file_exists(assets_destination_dir):
-    file_io.recursive_create_dir(assets_destination_dir)
 
+def _export_model_json_and_variables(model, saved_model_path):
+  """Save model variables and json structure into SavedModel subdirectories."""
+  # Save model configuration as a json string under assets folder.
+  model_json = model.to_json()
   model_json_filepath = os.path.join(
-      compat.as_bytes(assets_destination_dir),
-      compat.as_bytes(constants.SAVED_MODEL_FILENAME_JSON))
+      saved_model_utils.get_or_create_assets_dir(saved_model_path),
+      compat.as_text(constants.SAVED_MODEL_FILENAME_JSON))
   file_io.write_string_to_file(model_json_filepath, model_json)
 
-  # save model weights in checkpoint format.
-  checkpoint_destination_dir = os.path.join(
-      compat.as_bytes(saved_model_path),
-      compat.as_bytes(constants.VARIABLES_DIRECTORY))
+  # Save model weights in checkpoint format under variables folder.
+  saved_model_utils.get_or_create_variables_dir(saved_model_path)
+  checkpoint_prefix = saved_model_utils.get_variables_path(saved_model_path)
+  model.save_weights(checkpoint_prefix, save_format='tf', overwrite=True)
+  return checkpoint_prefix
 
-  if not file_io.file_exists(checkpoint_destination_dir):
-    file_io.recursive_create_dir(checkpoint_destination_dir)
 
-  checkpoint_prefix = os.path.join(
-      compat.as_text(checkpoint_destination_dir),
-      compat.as_text(constants.VARIABLES_FILENAME))
-  model.save_weights(checkpoint_prefix, save_format='tf', overwrite=True)
+def _get_var_list(model):
+  """Return list of all checkpointed saveable objects in the model."""
+  return checkpointable_utils.named_saveables(model)
+
+
+def _export_mode(
+    mode, has_saved_vars, builder, model, custom_objects, checkpoint_path):
+  """Export a model, and optionally save new vars from the clone model.
+
+  Args:
+    mode: A `tf.estimator.ModeKeys` string.
+    has_saved_vars: A `boolean` indicating whether the SavedModel has already
+      exported variables.
+    builder: A `SavedModelBuilder` object.
+    model: A `tf.keras.Model` object.
+    custom_objects: A dictionary mapping string names to custom classes
+      or functions.
+    checkpoint_path: String path to checkpoint.
+
+  Raises:
+    ValueError: If the train/eval mode is being exported, but the model does
+      not have an optimizer.
+  """
+  compile_clone = (mode != model_fn_lib.ModeKeys.PREDICT)
+  if compile_clone and not model.optimizer:
+    raise ValueError(
+        'Model does not have an optimizer. Cannot export mode %s' % mode)
+
+  model_graph = ops.get_default_graph()
+  with ops.Graph().as_default() as g:
+
+    K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN)
+
+    # Clone the model into blank graph. This will create placeholders for inputs
+    # and targets.
+    clone = models_lib.clone_and_build_model(
+        model, custom_objects=custom_objects, compile_clone=compile_clone)
+
+    # Make sure that iterations variable is added to the global step collection,
+    # to ensure that, when the SavedModel graph is loaded, the iterations
+    # variable is returned by `tf.train.get_global_step()`. This is required for
+    # compatibility with the SavedModelEstimator.
+    if compile_clone:
+      g.add_to_collection(ops.GraphKeys.GLOBAL_STEP, clone.optimizer.iterations)
+
+    # Extract update and train ops from train/test/predict functions.
+    if mode == model_fn_lib.ModeKeys.TRAIN:
+      clone._make_train_function()
+      builder._add_train_op(clone.train_function.updates_op)
+    elif mode == model_fn_lib.ModeKeys.EVAL:
+      clone._make_test_function()
+    else:
+      clone._make_predict_function()
+    g.get_collection_ref(ops.GraphKeys.UPDATE_OPS).extend(clone.state_updates)
+
+    clone_var_list = checkpointable_utils.named_saveables(clone)
+
+    with session.Session().as_default():
+      if has_saved_vars:
+        # Confirm all variables in the clone have an entry in the checkpoint.
+        status = clone.load_weights(checkpoint_path)
+        status.assert_existing_objects_matched()
+      else:
+        # Confirm that variables between the clone and model match up exactly,
+        # not counting optimizer objects. Optimizer objects are ignored because
+        # if the model has not trained, the slot variables will not have been
+        # created yet.
+        # TODO(b/113179535): Replace with checkpointable equivalence.
+        _assert_same_non_optimizer_objects(model, model_graph, clone, g)
+
+        # TODO(b/113178242): Use value transfer for checkpointable objects.
+        clone.load_weights(checkpoint_path)
+
+        # Add graph and variables to SavedModel.
+        # TODO(b/113134168): Switch to add_meta_graph_and_variables.
+        clone.save_weights(checkpoint_path, save_format='tf', overwrite=True)
+        builder._has_saved_variables = True
+
+    # Add graph to the SavedModel builder.
+    builder.add_meta_graph(
+        model_fn_lib.EXPORT_TAG_MAP[mode],
+        signature_def_map=_create_signature_def_map(clone, mode),
+        saver=saver_lib.Saver(clone_var_list),
+        main_op=variables.local_variables_initializer())
+    return None
+
+
+def _create_signature_def_map(model, mode):
+  """Create a SignatureDef map from a Keras model."""
+  inputs_dict = {name: x for name, x in zip(model.input_names, model.inputs)}
+  if model.optimizer:
+    targets_dict = {x.name.split(':')[0]: x
+                    for x in model.targets if x is not None}
+    inputs_dict.update(targets_dict)
+  outputs_dict = {name: x
+                  for name, x in zip(model.output_names, model.outputs)}
+  export_outputs = model_fn_lib.export_outputs_for_mode(
+      mode,
+      predictions=outputs_dict,
+      loss=model.total_loss if model.optimizer else None,
+      metrics=estimator_keras_util._convert_keras_metrics_to_estimator(model))
+  return export_helpers.build_all_signature_defs(
+      inputs_dict,
+      export_outputs=export_outputs,
+      serving_only=(mode == model_fn_lib.ModeKeys.PREDICT))
+
+
+def _assert_same_non_optimizer_objects(model, model_graph, clone, clone_graph):
+  """Assert model and clone contain the same checkpointable objects."""
+
+  def get_non_optimizer_objects(m, g):
+    """Gather set of model and optimizer checkpointable objects."""
+    # Set default graph because optimizer.variables() returns optimizer
+    # variables defined in the default graph.
+    with g.as_default():
+      all_objects = set(checkpointable_utils.list_objects(m))
+      optimizer_and_variables = set()
+      for obj in all_objects:
+        if isinstance(obj, optimizers.TFOptimizer):
+          optimizer_and_variables.update(checkpointable_utils.list_objects(obj))
+          optimizer_and_variables.update(set(obj.optimizer.variables()))
+      return all_objects - optimizer_and_variables
+
+  model_objects = get_non_optimizer_objects(model, model_graph)
+  clone_objects = get_non_optimizer_objects(clone, clone_graph)
+
+  if len(model_objects) != len(clone_objects):
+    raise errors.InternalError(
+        None, None,
+        'Model and clone must use the same variables.'
+        '\n\tModel variables: %s\n\t Clone variables: %s'
+        % (model_objects, clone_objects))
 
 
-def load_model(saved_model_path):
+def load_keras_model(saved_model_path):
   """Load a keras.Model from SavedModel.
 
   load_model reinstantiates model state by:
diff --git a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py
index 107ae1b07b..8a0dbef788 100644
--- a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py
+++ b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py
@@ -20,18 +20,35 @@ from __future__ import print_function
 
 import os
 import shutil
+
+from absl.testing import parameterized
 import numpy as np
 
 from tensorflow.contrib.saved_model.python.saved_model import keras_saved_model
 from tensorflow.python import keras
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+from tensorflow.python.estimator import model_fn as model_fn_lib
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.keras.engine import training
+from tensorflow.python.keras.utils import tf_utils
+from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import test
+from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import loader_impl
+from tensorflow.python.saved_model import signature_constants
 from tensorflow.python.training import training as training_module
 
 
 class TestModelSavingandLoading(test.TestCase):
 
+  def _save_model_dir(self, dirname='saved_model'):
+    temp_dir = self.get_temp_dir()
+    self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
+    return os.path.join(temp_dir, dirname)
+
   def test_saving_sequential_model(self):
     with self.test_session():
       model = keras.models.Sequential()
@@ -48,13 +65,11 @@ class TestModelSavingandLoading(test.TestCase):
       model.train_on_batch(x, y)
 
       ref_y = model.predict(x)
-      temp_dir = self.get_temp_dir()
-      self.addCleanup(shutil.rmtree, temp_dir)
 
-      temp_saved_model = os.path.join(temp_dir, 'saved_model')
-      keras_saved_model.save_model(model, temp_saved_model)
+      temp_saved_model = self._save_model_dir()
+      output_path = keras_saved_model.save_keras_model(model, temp_saved_model)
 
-      loaded_model = keras_saved_model.load_model(temp_saved_model)
+      loaded_model = keras_saved_model.load_keras_model(output_path)
       y = loaded_model.predict(x)
       self.assertAllClose(ref_y, y, atol=1e-05)
 
@@ -69,12 +84,9 @@ class TestModelSavingandLoading(test.TestCase):
       x = np.random.random((1, 3))
       ref_y = model.predict(x)
 
-      temp_dir = self.get_temp_dir()
-      self.addCleanup(shutil.rmtree, temp_dir)
-
-      temp_saved_model = os.path.join(temp_dir, 'saved_model')
-      keras_saved_model.save_model(model, temp_saved_model)
-      loaded_model = keras_saved_model.load_model(temp_saved_model)
+      temp_saved_model = self._save_model_dir()
+      output_path = keras_saved_model.save_keras_model(model, temp_saved_model)
+      loaded_model = keras_saved_model.load_keras_model(output_path)
 
       y = loaded_model.predict(x)
       self.assertAllClose(ref_y, y, atol=1e-05)
@@ -95,12 +107,10 @@ class TestModelSavingandLoading(test.TestCase):
       model.train_on_batch(x, y)
 
       ref_y = model.predict(x)
-      temp_dir = self.get_temp_dir()
-      self.addCleanup(shutil.rmtree, temp_dir)
 
-      temp_saved_model = os.path.join(temp_dir, 'saved_model')
-      keras_saved_model.save_model(model, temp_saved_model)
-      loaded_model = keras_saved_model.load_model(temp_saved_model)
+      temp_saved_model = self._save_model_dir()
+      output_path = keras_saved_model.save_keras_model(model, temp_saved_model)
+      loaded_model = keras_saved_model.load_keras_model(output_path)
 
       y = loaded_model.predict(x)
       self.assertAllClose(ref_y, y, atol=1e-05)
@@ -118,12 +128,10 @@ class TestModelSavingandLoading(test.TestCase):
       y = np.random.random((1, 3))
 
       ref_y = model.predict(x)
-      temp_dir = self.get_temp_dir()
-      self.addCleanup(shutil.rmtree, temp_dir)
 
-      temp_saved_model = os.path.join(temp_dir, 'saved_model')
-      keras_saved_model.save_model(model, temp_saved_model)
-      loaded_model = keras_saved_model.load_model(temp_saved_model)
+      temp_saved_model = self._save_model_dir()
+      output_path = keras_saved_model.save_keras_model(model, temp_saved_model)
+      loaded_model = keras_saved_model.load_keras_model(output_path)
 
       y = loaded_model.predict(x)
       self.assertAllClose(ref_y, y, atol=1e-05)
@@ -142,14 +150,13 @@ class TestModelSavingandLoading(test.TestCase):
       x = np.random.random((1, 3))
       y = np.random.random((1, 3))
       model.train_on_batch(x, y)
+      model.train_on_batch(x, y)
 
       ref_y = model.predict(x)
-      temp_dir = self.get_temp_dir()
-      self.addCleanup(shutil.rmtree, temp_dir)
 
-      temp_saved_model = os.path.join(temp_dir, 'saved_model')
-      keras_saved_model.save_model(model, temp_saved_model)
-      loaded_model = keras_saved_model.load_model(temp_saved_model)
+      temp_saved_model = self._save_model_dir()
+      output_path = keras_saved_model.save_keras_model(model, temp_saved_model)
+      loaded_model = keras_saved_model.load_keras_model(output_path)
       loaded_model.compile(
           loss='mse',
           optimizer=training_module.RMSPropOptimizer(0.1),
@@ -170,8 +177,10 @@ class TestModelSavingandLoading(test.TestCase):
       self.assertAllClose(ref_y, y, atol=1e-05)
 
       # test saving/loading again
-      keras_saved_model.save_model(loaded_model, temp_saved_model)
-      loaded_model = keras_saved_model.load_model(temp_saved_model)
+      temp_saved_model2 = self._save_model_dir('saved_model_2')
+      output_path2 = keras_saved_model.save_keras_model(
+          loaded_model, temp_saved_model2)
+      loaded_model = keras_saved_model.load_keras_model(output_path2)
       y = loaded_model.predict(x)
       self.assertAllClose(ref_y, y, atol=1e-05)
 
@@ -190,11 +199,231 @@ class TestModelSavingandLoading(test.TestCase):
         return self.layer2(self.layer1(inp))
 
     model = SubclassedModel()
-    temp_dir = self.get_temp_dir()
-    self.addCleanup(shutil.rmtree, temp_dir)
-    temp_saved_model = os.path.join(temp_dir, 'saved_model')
+
+    temp_saved_model = self._save_model_dir()
     with self.assertRaises(NotImplementedError):
-      keras_saved_model.save_model(model, temp_saved_model)
+      keras_saved_model.save_keras_model(model, temp_saved_model)
+
+
+class LayerWithLearningPhase(keras.engine.base_layer.Layer):
+
+  def call(self, x):
+    phase = keras.backend.learning_phase()
+    output = tf_utils.smart_cond(
+        phase, lambda: x * 0, lambda: array_ops.identity(x))
+    if not context.executing_eagerly():
+      output._uses_learning_phase = True  # pylint: disable=protected-access
+    return output
+
+  def compute_output_shape(self, input_shape):
+    return input_shape
+
+
+def functional_model(uses_learning_phase):
+  inputs = keras.layers.Input(shape=(3,))
+  x = keras.layers.Dense(2)(inputs)
+  x = keras.layers.Dense(3)(x)
+  if uses_learning_phase:
+    x = LayerWithLearningPhase()(x)
+  return keras.models.Model(inputs, x)
+
+
+def sequential_model(uses_learning_phase):
+  model = keras.models.Sequential()
+  model.add(keras.layers.Dense(2, input_shape=(3,)))
+  model.add(keras.layers.Dense(3))
+  if uses_learning_phase:
+    model.add(LayerWithLearningPhase())
+  return model
+
+
+def load_model(sess, path, mode):
+  tags = model_fn_lib.EXPORT_TAG_MAP[mode]
+  sig_def_key = (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+                 if mode == model_fn_lib.ModeKeys.PREDICT else mode)
+  meta_graph_def = loader_impl.load(sess, tags, path)
+  inputs = {
+      k: sess.graph.get_tensor_by_name(v.name)
+      for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()}
+  outputs = {
+      k: sess.graph.get_tensor_by_name(v.name)
+      for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()}
+  return inputs, outputs
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class TestModelSavedModelExport(test.TestCase, parameterized.TestCase):
+
+  def _save_model_dir(self, dirname='saved_model'):
+    temp_dir = self.get_temp_dir()
+    self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
+    return os.path.join(temp_dir, dirname)
+
+  @parameterized.parameters(
+      (functional_model, True, training_module.AdadeltaOptimizer(), True),
+      (functional_model, True, training_module.AdadeltaOptimizer(), False),
+      (functional_model, False, None, False),
+      (sequential_model, True, training_module.AdadeltaOptimizer(), True),
+      (sequential_model, True, training_module.AdadeltaOptimizer(), False),
+      (sequential_model, False, None, False))
+  def testSaveAndLoadSavedModelExport(
+      self, model_builder, uses_learning_phase, optimizer, train_before_export):
+    saved_model_path = self._save_model_dir()
+    with self.test_session(graph=ops.Graph()):
+      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
+      output_path = keras_saved_model.save_keras_model(model, saved_model_path)
+
+    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, output_path,
+                                   model_fn_lib.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, output_path,
+                                     model_fn_lib.ModeKeys.EVAL)
+
+        eval_results = sess.run(outputs, {inputs[input_name]: input_arr,
+                                          inputs[target_name]: target_arr})
+
+        self.assertEqual(int(train_before_export),
+                         sess.run(training_module.get_global_step()))
+        self.assertAllClose(ref_loss, eval_results['loss'], atol=1e-05)
+        self.assertAllClose(
+            ref_mae, eval_results['metrics/mae/update_op'], atol=1e-05)
+        self.assertAllClose(
+            ref_predict, eval_results['predictions/' + output_name], 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 = load_model(sess, output_path,
+                                     model_fn_lib.ModeKeys.TRAIN)
+        self.assertEqual(int(train_before_export),
+                         sess.run(training_module.get_global_step()))
+        self.assertIn('loss', outputs)
+        self.assertIn('metrics/mae/update_op', outputs)
+        self.assertIn('metrics/mae/value', outputs)
+        self.assertIn('predictions/' + output_name, outputs)
+
+        # Train for a step
+        train_op = ops.get_collection(constants.TRAIN_OP_KEY)
+        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 testSaveAndLoadSavedModelWithCustomObject(self):
+    saved_model_path = 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)
+      output_path = keras_saved_model.save_keras_model(
+          model, saved_model_path, custom_objects={'relu6': relu6})
+    with session.Session(graph=ops.Graph()) as sess:
+      inputs, outputs = load_model(sess, output_path,
+                                   model_fn_lib.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 testAssertModelCloneSameObjectsIgnoreOptimizer(self):
+    input_arr = np.random.random((1, 3))
+    target_arr = np.random.random((1, 3))
+
+    model_graph = ops.Graph()
+    clone_graph = ops.Graph()
+
+    # Create two models with the same layers but different optimizers.
+    with session.Session(graph=model_graph):
+      inputs = keras.layers.Input(shape=(3,))
+      x = keras.layers.Dense(2)(inputs)
+      x = keras.layers.Dense(3)(x)
+      model = keras.models.Model(inputs, x)
+
+      model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer())
+      model.train_on_batch(input_arr, target_arr)
+
+    with session.Session(graph=clone_graph):
+      inputs = keras.layers.Input(shape=(3,))
+      x = keras.layers.Dense(2)(inputs)
+      x = keras.layers.Dense(3)(x)
+      clone = keras.models.Model(inputs, x)
+      clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001))
+      clone.train_on_batch(input_arr, target_arr)
+
+    keras_saved_model._assert_same_non_optimizer_objects(
+        model, model_graph, clone, clone_graph)
+
+  def testAssertModelCloneSameObjectsThrowError(self):
+    input_arr = np.random.random((1, 3))
+    target_arr = np.random.random((1, 3))
+
+    model_graph = ops.Graph()
+    clone_graph = ops.Graph()
+
+    # Create two models with the same layers but different optimizers.
+    with session.Session(graph=model_graph):
+      inputs = keras.layers.Input(shape=(3,))
+      x = keras.layers.Dense(2)(inputs)
+      x = keras.layers.Dense(3)(x)
+      model = keras.models.Model(inputs, x)
+
+      model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer())
+      model.train_on_batch(input_arr, target_arr)
+
+    with session.Session(graph=clone_graph):
+      inputs = keras.layers.Input(shape=(3,))
+      x = keras.layers.Dense(2)(inputs)
+      x = keras.layers.Dense(4)(x)
+      x = keras.layers.Dense(3)(x)
+      clone = keras.models.Model(inputs, x)
+      clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001))
+      clone.train_on_batch(input_arr, target_arr)
+
+    with self.assertRaisesRegexp(
+        errors.InternalError, 'Model and clone must use the same variables.'):
+      keras_saved_model._assert_same_non_optimizer_objects(
+          model, model_graph, clone, clone_graph)
 
 
 if __name__ == '__main__':