Automated rollback of commit 069f808

PiperOrigin-RevId: 210656847

Author: k-w-w

tensorflow/contrib/saved_model/BUILD

@@ -36,7 +36,6 @@ py_library(
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
     deps = [
-        ":keras_saved_model",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:framework_ops",
         "//tensorflow/python:lib",
@@ -102,33 +101,23 @@ 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",
+        "//tensorflow/python/saved_model:constants",
     ],
 )
 
 py_test(
     name = "keras_saved_model_test",
-    size = "medium",
+    size = "small",
     srcs = ["python/saved_model/keras_saved_model_test.py"],
     srcs_version = "PY2AND3",
     deps = [
-        ":keras_saved_model",
+        ":saved_model_py",
         "//tensorflow/python:client_testlib",
         "//tensorflow/python:training",
         "//tensorflow/python/keras",
         "//third_party/py/numpy",
-        "@absl_py//absl/testing:parameterized",
     ],
 )

tensorflow/contrib/saved_model/__init__.py

@@ -26,13 +26,10 @@
 # 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,wildcard-import,line-too-long
+# pylint: enable=unused-import,widcard-import,line-too-long
 
 from tensorflow.python.util.all_util import remove_undocumented
 
-_allowed_symbols = [
-    "get_signature_def_by_key",
-    "load_keras_model",
-    "save_keras_model"]
+_allowed_symbols = ["get_signature_def_by_key", "load_model", "save_model"]
 
 remove_undocumented(__name__, _allowed_symbols)

tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py

@@ -20,270 +20,64 @@
 
 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_keras_model(
-    model, saved_model_path, custom_objects=None, as_text=None):
+def save_model(model, saved_model_path):
   """Save a `tf.keras.Model` into Tensorflow SavedModel format.
 
-  `save_model` generates new files/folders under the `saved_model_path` folder:
+  `save_model` generates such 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.
 
-  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.
+  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`.
 
   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.
   """
   if not model._is_graph_network:
     raise NotImplementedError
 
-  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
+  # 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))
 
+  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(
-      saved_model_utils.get_or_create_assets_dir(saved_model_path),
-      compat.as_text(constants.SAVED_MODEL_FILENAME_JSON))
+      compat.as_bytes(assets_destination_dir),
+      compat.as_bytes(constants.SAVED_MODEL_FILENAME_JSON))
   file_io.write_string_to_file(model_json_filepath, model_json)
 
-  # 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
-
-
-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
+  # save model weights in checkpoint format.
+  checkpoint_destination_dir = os.path.join(
+      compat.as_bytes(saved_model_path),
+      compat.as_bytes(constants.VARIABLES_DIRECTORY))
 
-  model_objects = get_non_optimizer_objects(model, model_graph)
-  clone_objects = get_non_optimizer_objects(clone, clone_graph)
+  if not file_io.file_exists(checkpoint_destination_dir):
+    file_io.recursive_create_dir(checkpoint_destination_dir)
 
-  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))
+  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 load_keras_model(saved_model_path):
+def load_model(saved_model_path):
   """Load a keras.Model from SavedModel.
 
   load_model reinstantiates model state by: