chore: import upstream snapshot with attribution

tensorflow/python/keras/saving/BUILD

@@ -0,0 +1,61 @@
+# Description:
+#   Contains the Keras save model API (internal TensorFlow version).
+
+load("@xla//third_party/rules_python/python:defs.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    # TODO(scottzhu): Remove non-keras deps from TF.
+    default_visibility = [
+        "//tensorflow/python/distribute:__pkg__",
+        "//tensorflow/python/keras:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+filegroup(
+    name = "all_py_srcs",
+    srcs = glob([
+        "*.py",
+        "**/*.py",
+    ]),
+    visibility = ["//tensorflow/python/keras/google/private_tf_api_test:__pkg__"],
+)
+
+py_library(
+    name = "saving",
+    srcs = [
+        "__init__.py",
+        "hdf5_format.py",
+        "model_config.py",
+        "save.py",
+        "saved_model_experimental.py",
+        "saving_utils.py",
+    ],
+    srcs_version = "PY3",
+    strict_deps = False,
+    deps = [
+        "//tensorflow/python/checkpoint:graph_view",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/keras:backend",
+        "//tensorflow/python/keras:losses",
+        "//tensorflow/python/keras:optimizers",
+        "//tensorflow/python/keras:regularizers",
+        "//tensorflow/python/keras/engine:input_spec",
+        "//tensorflow/python/keras/mixed_precision:autocast_variable",
+        "//tensorflow/python/keras/protobuf:saved_metadata_proto_py",
+        "//tensorflow/python/keras/saving/saved_model",
+        "//tensorflow/python/keras/saving/utils_v1",
+        "//tensorflow/python/keras/utils:engine_utils",
+        "//tensorflow/python/keras/utils:metrics_utils",
+        "//tensorflow/python/keras/utils:mode_keys",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model",
+        "//tensorflow/python/training:saver",
+        "@pypi//h5py",
+    ],
+)

tensorflow/python/keras/saving/hdf5_format.py

@@ -0,0 +1,902 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# pylint: disable=protected-access
+"""Functions for saving and loading a Keras Model from HDF5 format."""
+
+import json
+import os
+
+import numpy as np
+
+from tensorflow.python.keras import backend
+from tensorflow.python.keras import optimizer_v1
+from tensorflow.python.keras.saving import model_config as model_config_lib
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import json_utils
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.keras.utils.io_utils import ask_to_proceed_with_overwrite
+from tensorflow.python.ops import variables as variables_module
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import tf_logging as logging
+
+
+# pylint: disable=g-import-not-at-top
+try:
+  # Disable loading HDF5 plugins from a default path and prevent ZDI-CAN-25480.
+  # Importing h5py prior to importing tensorflow will restore the old behavior.
+  os.environ.setdefault('HDF5_PLUGIN_PATH', os.devnull)
+  import h5py
+  HDF5_OBJECT_HEADER_LIMIT = 64512
+except ImportError:
+  h5py = None
+# pylint: enable=g-import-not-at-top
+
+# TODO(b/134426265): Switch back to single-quotes to match the rest of the file
+# once the issue with copybara is fixed.
+# pylint:disable=g-inconsistent-quotes
+sequential_lib = LazyLoader(
+    "sequential_lib", globals(),
+    "tensorflow.python.keras.engine.sequential")
+# pylint:enable=g-inconsistent-quotes
+
+
+def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True):
+  """Saves a model to a HDF5 file.
+
+  The saved model contains:
+      - the model's configuration (topology)
+      - the model's weights
+      - the model's optimizer's state (if any)
+
+  Thus the saved model can be reinstantiated in
+  the exact same state, without any of the code
+  used for model definition or training.
+
+  Args:
+      model: Keras model instance to be saved.
+      filepath: One of the following:
+          - String, path where to save the model
+          - `h5py.File` object where to save the model
+      overwrite: Whether we should overwrite any existing
+          model at the target location, or instead
+          ask the user with a manual prompt.
+      include_optimizer: If True, save optimizer's state together.
+
+  Raises:
+      ImportError: if h5py is not available.
+  """
+
+  if h5py is None:
+    raise ImportError('`save_model` requires h5py.')
+
+  # TODO(psv) Add warning when we save models that contain non-serializable
+  # entities like metrics added using `add_metric` and losses added using
+  # `add_loss.`
+  if len(model.weights) != len(model._undeduplicated_weights):
+    logging.warning('Found duplicated `Variable`s in Model\'s `weights`. '
+                    'This is usually caused by `Variable`s being shared by '
+                    'Layers in the Model. These `Variable`s will be treated '
+                    'as separate `Variable`s when the Model is restored. To '
+                    'avoid this, please save with `save_format="tf"`.')
+
+  if not isinstance(filepath, h5py.File):
+    # If file exists and should not be overwritten.
+    if not overwrite and os.path.isfile(filepath):
+      proceed = ask_to_proceed_with_overwrite(filepath)
+      if not proceed:
+        return
+
+    # Try creating dir if not exist
+    dirpath = os.path.dirname(filepath)
+    if not os.path.exists(dirpath):
+      gfile.MakeDirs(dirpath)
+
+    f = h5py.File(filepath, mode='w')
+    opened_new_file = True
+  else:
+    f = filepath
+    opened_new_file = False
+
+  try:
+    model_metadata = saving_utils.model_metadata(model, include_optimizer)
+    for k, v in model_metadata.items():
+      if isinstance(v, (dict, list, tuple)):
+        f.attrs[k] = json.dumps(
+            v, default=json_utils.get_json_type).encode('utf8')
+      else:
+        f.attrs[k] = v
+
+    model_weights_group = f.create_group('model_weights')
+    model_layers = model.layers
+    save_weights_to_hdf5_group(model_weights_group, model_layers)
+
+    # TODO(b/128683857): Add integration tests between tf.keras and external
+    # Keras, to avoid breaking TF.js users.
+    if (include_optimizer and model.optimizer and
+        not isinstance(model.optimizer, optimizer_v1.TFOptimizer)):
+      save_optimizer_weights_to_hdf5_group(f, model.optimizer)
+
+    f.flush()
+  finally:
+    if opened_new_file:
+      f.close()
+
+
+def load_model_from_hdf5(filepath, custom_objects=None, compile=True):  # pylint: disable=redefined-builtin
+  """Loads a model saved via `save_model_to_hdf5`.
+
+  Args:
+      filepath: One of the following:
+          - String, path to the saved model
+          - `h5py.File` object from which to load the model
+      custom_objects: Optional dictionary mapping names
+          (strings) to custom classes or functions to be
+          considered during deserialization.
+      compile: Boolean, whether to compile the model
+          after loading.
+
+  Returns:
+      A Keras model instance. If an optimizer was found
+      as part of the saved model, the model is already
+      compiled. Otherwise, the model is uncompiled and
+      a warning will be displayed. When `compile` is set
+      to False, the compilation is omitted without any
+      warning.
+
+  Raises:
+      ImportError: if h5py is not available.
+      ValueError: In case of an invalid savefile.
+  """
+  if h5py is None:
+    raise ImportError('`load_model` requires h5py.')
+
+  if not custom_objects:
+    custom_objects = {}
+
+  opened_new_file = not isinstance(filepath, h5py.File)
+  if opened_new_file:
+    f = h5py.File(filepath, mode='r')
+  else:
+    f = filepath
+
+  model = None
+  try:
+    # instantiate model
+    model_config = f.attrs.get('model_config')
+    if model_config is None:
+      raise ValueError('No model found in config file.')
+    if hasattr(model_config, 'decode'):
+      model_config = model_config.decode('utf-8')
+    model_config = json_utils.decode(model_config)
+    model = model_config_lib.model_from_config(model_config,
+                                               custom_objects=custom_objects)
+
+    # set weights
+    load_weights_from_hdf5_group(f['model_weights'], model.layers)
+
+    if compile:
+      # instantiate optimizer
+      training_config = f.attrs.get('training_config')
+      if hasattr(training_config, 'decode'):
+        training_config = training_config.decode('utf-8')
+      if training_config is None:
+        logging.warning('No training configuration found in the save file, so '
+                        'the model was *not* compiled. Compile it manually.')
+        return model
+      training_config = json_utils.decode(training_config)
+
+      # Compile model.
+      model.compile(**saving_utils.compile_args_from_training_config(
+          training_config, custom_objects), from_serialized=True)
+      saving_utils.try_build_compiled_arguments(model)
+
+      # Set optimizer weights.
+      if 'optimizer_weights' in f:
+        try:
+          model.optimizer._create_all_weights(model.trainable_variables)
+        except (NotImplementedError, AttributeError):
+          logging.warning(
+              'Error when creating the weights of optimizer {}, making it '
+              'impossible to restore the saved optimizer state. As a result, '
+              'your model is starting with a freshly initialized optimizer.')
+
+        optimizer_weight_values = load_optimizer_weights_from_hdf5_group(f)
+        try:
+          model.optimizer.set_weights(optimizer_weight_values)
+        except ValueError:
+          logging.warning('Error in loading the saved optimizer '
+                          'state. As a result, your model is '
+                          'starting with a freshly initialized '
+                          'optimizer.')
+  finally:
+    if opened_new_file:
+      f.close()
+  return model
+
+
+def preprocess_weights_for_loading(layer,
+                                   weights,
+                                   original_keras_version=None,
+                                   original_backend=None):
+  """Preprocess layer weights between different Keras formats.
+
+  Converts layers weights from Keras 1 format to Keras 2 and also weights of
+  CuDNN layers in Keras 2.
+
+  Args:
+      layer: Layer instance.
+      weights: List of weights values (Numpy arrays).
+      original_keras_version: Keras version for the weights, as a string.
+      original_backend: Keras backend the weights were trained with,
+          as a string.
+
+  Returns:
+      A list of weights values (Numpy arrays).
+  """
+  def convert_nested_bidirectional(weights):
+    """Converts layers nested in `Bidirectional` wrapper.
+
+    This function uses `preprocess_weights_for_loading()` for converting
+    layers.
+
+    Args:
+        weights: List of weights values (Numpy arrays).
+
+    Returns:
+        A list of weights values (Numpy arrays).
+    """
+    num_weights_per_layer = len(weights) // 2
+    forward_weights = preprocess_weights_for_loading(
+        layer.forward_layer, weights[:num_weights_per_layer],
+        original_keras_version, original_backend)
+    backward_weights = preprocess_weights_for_loading(
+        layer.backward_layer, weights[num_weights_per_layer:],
+        original_keras_version, original_backend)
+    return forward_weights + backward_weights
+
+  def convert_nested_time_distributed(weights):
+    """Converts layers nested in `TimeDistributed` wrapper.
+
+    This function uses `preprocess_weights_for_loading()` for converting nested
+    layers.
+
+    Args:
+        weights: List of weights values (Numpy arrays).
+
+    Returns:
+        A list of weights values (Numpy arrays).
+    """
+    return preprocess_weights_for_loading(
+        layer.layer, weights, original_keras_version, original_backend)
+
+  def convert_nested_model(weights):
+    """Converts layers nested in `Model` or `Sequential`.
+
+    This function uses `preprocess_weights_for_loading()` for converting nested
+    layers.
+
+    Args:
+        weights: List of weights values (Numpy arrays).
+
+    Returns:
+        A list of weights values (Numpy arrays).
+    """
+    trainable_weights = weights[:len(layer.trainable_weights)]
+    non_trainable_weights = weights[len(layer.trainable_weights):]
+
+    new_trainable_weights = []
+    new_non_trainable_weights = []
+
+    for sublayer in layer.layers:
+      num_trainable_weights = len(sublayer.trainable_weights)
+      num_non_trainable_weights = len(sublayer.non_trainable_weights)
+      if sublayer.weights:
+        preprocessed = preprocess_weights_for_loading(
+            layer=sublayer,
+            weights=(trainable_weights[:num_trainable_weights] +
+                     non_trainable_weights[:num_non_trainable_weights]),
+            original_keras_version=original_keras_version,
+            original_backend=original_backend)
+        new_trainable_weights.extend(preprocessed[:num_trainable_weights])
+        new_non_trainable_weights.extend(preprocessed[num_trainable_weights:])
+
+        trainable_weights = trainable_weights[num_trainable_weights:]
+        non_trainable_weights = non_trainable_weights[
+            num_non_trainable_weights:]
+
+    return new_trainable_weights + new_non_trainable_weights
+
+  # Convert layers nested in Bidirectional/Model/Sequential.
+  # Both transformation should be ran for both Keras 1->2 conversion
+  # and for conversion of CuDNN layers.
+  if layer.__class__.__name__ == 'Bidirectional':
+    weights = convert_nested_bidirectional(weights)
+  if layer.__class__.__name__ == 'TimeDistributed':
+    weights = convert_nested_time_distributed(weights)
+  elif layer.__class__.__name__ in ['Model', 'Sequential', 'Functional']:
+    weights = convert_nested_model(weights)
+
+  if original_keras_version == '1':
+    if layer.__class__.__name__ == 'TimeDistributed':
+      weights = preprocess_weights_for_loading(
+          layer.layer, weights, original_keras_version, original_backend)
+
+    if layer.__class__.__name__ == 'Conv1D':
+      shape = weights[0].shape
+      # Handle Keras 1.1 format
+      if shape[:2] != (layer.kernel_size[0], 1) or shape[3] != layer.filters:
+        # Legacy shape:
+        # (filters, input_dim, filter_length, 1)
+        assert shape[0] == layer.filters and shape[2:] == (layer.kernel_size[0],
+                                                           1)
+        weights[0] = np.transpose(weights[0], (2, 3, 1, 0))
+      weights[0] = weights[0][:, 0, :, :]
+
+    if layer.__class__.__name__ == 'Conv2D':
+      if layer.data_format == 'channels_first':
+        # old: (filters, stack_size, kernel_rows, kernel_cols)
+        # new: (kernel_rows, kernel_cols, stack_size, filters)
+        weights[0] = np.transpose(weights[0], (2, 3, 1, 0))
+
+    if layer.__class__.__name__ == 'Conv2DTranspose':
+      if layer.data_format == 'channels_last':
+        # old: (kernel_rows, kernel_cols, stack_size, filters)
+        # new: (kernel_rows, kernel_cols, filters, stack_size)
+        weights[0] = np.transpose(weights[0], (0, 1, 3, 2))
+      if layer.data_format == 'channels_first':
+        # old: (filters, stack_size, kernel_rows, kernel_cols)
+        # new: (kernel_rows, kernel_cols, filters, stack_size)
+        weights[0] = np.transpose(weights[0], (2, 3, 0, 1))
+
+    if layer.__class__.__name__ == 'Conv3D':
+      if layer.data_format == 'channels_first':
+        # old: (filters, stack_size, ...)
+        # new: (..., stack_size, filters)
+        weights[0] = np.transpose(weights[0], (2, 3, 4, 1, 0))
+
+    if layer.__class__.__name__ == 'GRU':
+      if len(weights) == 9:
+        kernel = np.concatenate([weights[0], weights[3], weights[6]], axis=-1)
+        recurrent_kernel = np.concatenate(
+            [weights[1], weights[4], weights[7]], axis=-1)
+        bias = np.concatenate([weights[2], weights[5], weights[8]], axis=-1)
+        weights = [kernel, recurrent_kernel, bias]
+
+    if layer.__class__.__name__ == 'LSTM':
+      if len(weights) == 12:
+        # old: i, c, f, o
+        # new: i, f, c, o
+        kernel = np.concatenate(
+            [weights[0], weights[6], weights[3], weights[9]], axis=-1)
+        recurrent_kernel = np.concatenate(
+            [weights[1], weights[7], weights[4], weights[10]], axis=-1)
+        bias = np.concatenate(
+            [weights[2], weights[8], weights[5], weights[11]], axis=-1)
+        weights = [kernel, recurrent_kernel, bias]
+
+    if layer.__class__.__name__ == 'ConvLSTM2D':
+      if len(weights) == 12:
+        kernel = np.concatenate(
+            [weights[0], weights[6], weights[3], weights[9]], axis=-1)
+        recurrent_kernel = np.concatenate(
+            [weights[1], weights[7], weights[4], weights[10]], axis=-1)
+        bias = np.concatenate(
+            [weights[2], weights[8], weights[5], weights[11]], axis=-1)
+        if layer.data_format == 'channels_first':
+          # old: (filters, stack_size, kernel_rows, kernel_cols)
+          # new: (kernel_rows, kernel_cols, stack_size, filters)
+          kernel = np.transpose(kernel, (2, 3, 1, 0))
+          recurrent_kernel = np.transpose(recurrent_kernel, (2, 3, 1, 0))
+        weights = [kernel, recurrent_kernel, bias]
+
+  conv_layers = ['Conv1D', 'Conv2D', 'Conv3D', 'Conv2DTranspose', 'ConvLSTM2D']
+  if layer.__class__.__name__ in conv_layers:
+    if backend.int_shape(layer.weights[0]) != weights[0].shape:
+      weights[0] = np.transpose(weights[0], (3, 2, 0, 1))
+      if layer.__class__.__name__ == 'ConvLSTM2D':
+        weights[1] = np.transpose(weights[1], (3, 2, 0, 1))
+
+  # convert CuDNN layers
+  return _convert_rnn_weights(layer, weights)
+
+
+def _convert_rnn_weights(layer, weights):
+  """Converts weights for RNN layers between native and CuDNN format.
+
+  Input kernels for each gate are transposed and converted between Fortran
+  and C layout, recurrent kernels are transposed. For LSTM biases are summed/
+  split in half, for GRU biases are reshaped.
+
+  Weights can be converted in both directions between `LSTM` and`CuDNNSLTM`
+  and between `CuDNNGRU` and `GRU(reset_after=True)`. Default `GRU` is not
+  compatible with `CuDNNGRU`.
+
+  For missing biases in `LSTM`/`GRU` (`use_bias=False`) no conversion is made.
+
+  Args:
+      layer: Target layer instance.
+      weights: List of source weights values (input kernels, recurrent
+          kernels, [biases]) (Numpy arrays).
+
+  Returns:
+      A list of converted weights values (Numpy arrays).
+
+  Raises:
+      ValueError: for incompatible GRU layer/weights or incompatible biases
+  """
+
+  def transform_kernels(kernels, func, n_gates):
+    """Transforms kernel for each gate separately using given function.
+
+    Args:
+        kernels: Stacked array of kernels for individual gates.
+        func: Function applied to kernel of each gate.
+        n_gates: Number of gates (4 for LSTM, 3 for GRU).
+
+    Returns:
+        Stacked array of transformed kernels.
+    """
+    return np.hstack([func(k) for k in np.hsplit(kernels, n_gates)])
+
+  def transpose_input(from_cudnn):
+    """Makes a function that transforms input kernels from/to CuDNN format.
+
+    It keeps the shape, but changes between the layout (Fortran/C). Eg.:
+
+    ```
+    Keras                 CuDNN
+    [[0, 1, 2],  <--->  [[0, 2, 4],
+     [3, 4, 5]]          [1, 3, 5]]
+    ```
+
+    It can be passed to `transform_kernels()`.
+
+    Args:
+        from_cudnn: `True` if source weights are in CuDNN format, `False`
+            if they're in plain Keras format.
+
+    Returns:
+        Function that converts input kernel to the other format.
+    """
+    order = 'F' if from_cudnn else 'C'
+
+    def transform(kernel):
+      return kernel.T.reshape(kernel.shape, order=order)
+
+    return transform
+
+  target_class = layer.__class__.__name__
+
+  # convert the weights between CuDNNLSTM and LSTM
+  if target_class in ['LSTM', 'CuDNNLSTM'] and len(weights) == 3:
+    # determine if we're loading a CuDNNLSTM layer
+    # from the number of bias weights:
+    # CuDNNLSTM has (units * 8) weights; while LSTM has (units * 4)
+    # if there's no bias weight in the file, skip this conversion
+    units = weights[1].shape[0]
+    bias_shape = weights[2].shape
+    n_gates = 4
+
+    if bias_shape == (2 * units * n_gates,):
+      source = 'CuDNNLSTM'
+    elif bias_shape == (units * n_gates,):
+      source = 'LSTM'
+    else:
+      raise ValueError('Invalid bias shape: ' + str(bias_shape))
+
+    def convert_lstm_weights(weights, from_cudnn=True):
+      """Converts the weights between CuDNNLSTM and LSTM.
+
+      Args:
+        weights: Original weights.
+        from_cudnn: Indicates whether original weights are from CuDNN layer.
+
+      Returns:
+        Updated weights compatible with LSTM.
+      """
+
+      # Transpose (and reshape) input and recurrent kernels
+      kernels = transform_kernels(weights[0], transpose_input(from_cudnn),
+                                  n_gates)
+      recurrent_kernels = transform_kernels(weights[1], lambda k: k.T, n_gates)
+      if from_cudnn:
+        # merge input and recurrent biases into a single set
+        biases = np.sum(np.split(weights[2], 2, axis=0), axis=0)
+      else:
+        # Split single set of biases evenly to two sets. The way of
+        # splitting doesn't matter as long as the two sets sum is kept.
+        biases = np.tile(0.5 * weights[2], 2)
+      return [kernels, recurrent_kernels, biases]
+
+    if source != target_class:
+      weights = convert_lstm_weights(weights, from_cudnn=source == 'CuDNNLSTM')
+
+  # convert the weights between CuDNNGRU and GRU(reset_after=True)
+  if target_class in ['GRU', 'CuDNNGRU'] and len(weights) == 3:
+    # We can determine the source of the weights from the shape of the bias.
+    # If there is no bias we skip the conversion since
+    # CuDNNGRU always has biases.
+
+    units = weights[1].shape[0]
+    bias_shape = weights[2].shape
+    n_gates = 3
+
+    def convert_gru_weights(weights, from_cudnn=True):
+      """Converts the weights between CuDNNGRU and GRU.
+
+      Args:
+        weights: Original weights.
+        from_cudnn: Indicates whether original weights are from CuDNN layer.
+
+      Returns:
+        Updated weights compatible with GRU.
+      """
+
+      kernels = transform_kernels(weights[0], transpose_input(from_cudnn),
+                                  n_gates)
+      recurrent_kernels = transform_kernels(weights[1], lambda k: k.T, n_gates)
+      biases = np.array(weights[2]).reshape((2, -1) if from_cudnn else -1)
+      return [kernels, recurrent_kernels, biases]
+
+    if bias_shape == (2 * units * n_gates,):
+      source = 'CuDNNGRU'
+    elif bias_shape == (2, units * n_gates):
+      source = 'GRU(reset_after=True)'
+    elif bias_shape == (units * n_gates,):
+      source = 'GRU(reset_after=False)'
+    else:
+      raise ValueError('Invalid bias shape: ' + str(bias_shape))
+
+    if target_class == 'CuDNNGRU':
+      target = 'CuDNNGRU'
+    elif layer.reset_after:
+      target = 'GRU(reset_after=True)'
+    else:
+      target = 'GRU(reset_after=False)'
+
+    # only convert between different types
+    if source != target:
+      types = (source, target)
+      if 'GRU(reset_after=False)' in types:
+        raise ValueError('%s is not compatible with %s' % types)
+      if source == 'CuDNNGRU':
+        weights = convert_gru_weights(weights, from_cudnn=True)
+      elif source == 'GRU(reset_after=True)':
+        weights = convert_gru_weights(weights, from_cudnn=False)
+
+  return weights
+
+
+def save_optimizer_weights_to_hdf5_group(hdf5_group, optimizer):
+  """Saves optimizer weights of a optimizer to a HDF5 group.
+
+  Args:
+      hdf5_group: HDF5 group.
+      optimizer: optimizer instance.
+  """
+
+  symbolic_weights = getattr(optimizer, 'weights')
+  if symbolic_weights:
+    weights_group = hdf5_group.create_group('optimizer_weights')
+    weight_names = [str(w.name).encode('utf8') for w in symbolic_weights]
+    save_attributes_to_hdf5_group(weights_group, 'weight_names', weight_names)
+    weight_values = backend.batch_get_value(symbolic_weights)
+    for name, val in zip(weight_names, weight_values):
+      param_dset = weights_group.create_dataset(
+          name, val.shape, dtype=val.dtype)
+      if not val.shape:
+        # scalar
+        param_dset[()] = val
+      else:
+        param_dset[:] = val
+
+
+def load_optimizer_weights_from_hdf5_group(hdf5_group):
+  """Load optimizer weights from a HDF5 group.
+
+  Args:
+      hdf5_group: A pointer to a HDF5 group.
+
+  Returns:
+      data: List of optimizer weight names.
+  """
+  weights_group = hdf5_group['optimizer_weights']
+  optimizer_weight_names = load_attributes_from_hdf5_group(
+      weights_group, 'weight_names')
+  return [weights_group[weight_name] for weight_name in optimizer_weight_names]
+
+
+def save_weights_to_hdf5_group(f, layers):
+  """Saves the weights of a list of layers to a HDF5 group.
+
+  Args:
+      f: HDF5 group.
+      layers: List of layer instances.
+  """
+  from tensorflow.python.keras import __version__ as keras_version  # pylint: disable=g-import-not-at-top
+
+  save_attributes_to_hdf5_group(
+      f, 'layer_names', [layer.name.encode('utf8') for layer in layers])
+  f.attrs['backend'] = backend.backend().encode('utf8')
+  f.attrs['keras_version'] = str(keras_version).encode('utf8')
+
+  # Sort model layers by layer name to ensure that group names are strictly
+  # growing to avoid prefix issues.
+  for layer in sorted(layers, key=lambda x: x.name):
+    g = f.create_group(layer.name)
+    weights = _legacy_weights(layer)
+    weight_values = backend.batch_get_value(weights)
+    weight_names = [w.name.encode('utf8') for w in weights]
+    save_attributes_to_hdf5_group(g, 'weight_names', weight_names)
+    for name, val in zip(weight_names, weight_values):
+      param_dset = g.create_dataset(name, val.shape, dtype=val.dtype)
+      if not val.shape:
+        # scalar
+        param_dset[()] = val
+      else:
+        param_dset[:] = val
+
+
+def load_weights_from_hdf5_group(f, layers):
+  """Implements topological (order-based) weight loading.
+
+  Args:
+      f: A pointer to a HDF5 group.
+      layers: a list of target layers.
+
+  Raises:
+      ValueError: in case of mismatch between provided layers
+          and weights file.
+  """
+  if 'keras_version' in f.attrs:
+    original_keras_version = f.attrs['keras_version']
+    if hasattr(original_keras_version, 'decode'):
+      original_keras_version = original_keras_version.decode('utf8')
+  else:
+    original_keras_version = '1'
+  if 'backend' in f.attrs:
+    original_backend = f.attrs['backend']
+    if hasattr(original_backend, 'decode'):
+      original_backend = original_backend.decode('utf8')
+  else:
+    original_backend = None
+
+  filtered_layers = []
+  for layer in layers:
+    weights = _legacy_weights(layer)
+    if weights:
+      filtered_layers.append(layer)
+
+  layer_names = load_attributes_from_hdf5_group(f, 'layer_names')
+  filtered_layer_names = []
+  for name in layer_names:
+    g = f[name]
+    weight_names = load_attributes_from_hdf5_group(g, 'weight_names')
+    if weight_names:
+      filtered_layer_names.append(name)
+  layer_names = filtered_layer_names
+  if len(layer_names) != len(filtered_layers):
+    raise ValueError('You are trying to load a weight file '
+                     'containing ' + str(len(layer_names)) +
+                     ' layers into a model with ' + str(len(filtered_layers)) +
+                     ' layers.')
+
+  # We batch weight value assignments in a single backend call
+  # which provides a speedup in TensorFlow.
+  weight_value_tuples = []
+  for k, name in enumerate(layer_names):
+    g = f[name]
+    weight_names = load_attributes_from_hdf5_group(g, 'weight_names')
+    weight_values = [np.asarray(g[weight_name]) for weight_name in weight_names]
+    layer = filtered_layers[k]
+    symbolic_weights = _legacy_weights(layer)
+    weight_values = preprocess_weights_for_loading(
+        layer, weight_values, original_keras_version, original_backend)
+    if len(weight_values) != len(symbolic_weights):
+      raise ValueError('Layer #' + str(k) + ' (named "' + layer.name +
+                       '" in the current model) was found to '
+                       'correspond to layer ' + name + ' in the save file. '
+                       'However the new layer ' + layer.name + ' expects ' +
+                       str(len(symbolic_weights)) +
+                       ' weights, but the saved weights have ' +
+                       str(len(weight_values)) + ' elements.')
+    weight_value_tuples += zip(symbolic_weights, weight_values)
+  backend.batch_set_value(weight_value_tuples)
+
+
+def load_weights_from_hdf5_group_by_name(
+    f, layers, skip_mismatch=False):
+  """Implements name-based weight loading.
+
+  (instead of topological weight loading).
+
+  Layers that have no matching name are skipped.
+
+  Args:
+      f: A pointer to a HDF5 group.
+      layers: a list of target layers.
+      skip_mismatch: Boolean, whether to skip loading of layers
+          where there is a mismatch in the number of weights,
+          or a mismatch in the shape of the weights.
+
+  Raises:
+      ValueError: in case of mismatch between provided layers
+          and weights file and skip_match=False.
+  """
+  if 'keras_version' in f.attrs:
+    original_keras_version = f.attrs['keras_version']
+    if hasattr(original_keras_version, 'decode'):
+      original_keras_version = original_keras_version.decode('utf8')
+  else:
+    original_keras_version = '1'
+  if 'backend' in f.attrs:
+    original_backend = f.attrs['backend']
+    if hasattr(original_backend, 'decode'):
+      original_backend = original_backend.decode('utf8')
+  else:
+    original_backend = None
+
+  # New file format.
+  layer_names = load_attributes_from_hdf5_group(f, 'layer_names')
+
+  # Reverse index of layer name to list of layers with name.
+  index = {}
+  for layer in layers:
+    if layer.name:
+      index.setdefault(layer.name, []).append(layer)
+
+  # We batch weight value assignments in a single backend call
+  # which provides a speedup in TensorFlow.
+  weight_value_tuples = []
+  for k, name in enumerate(layer_names):
+    g = f[name]
+    weight_names = load_attributes_from_hdf5_group(g, 'weight_names')
+    weight_values = [np.asarray(g[weight_name]) for weight_name in weight_names]
+
+    for layer in index.get(name, []):
+      symbolic_weights = _legacy_weights(layer)
+      weight_values = preprocess_weights_for_loading(
+          layer, weight_values, original_keras_version, original_backend)
+      if len(weight_values) != len(symbolic_weights):
+        if skip_mismatch:
+          logging.warning('Skipping loading of weights for '
+                          'layer {}'.format(layer.name) + ' due to mismatch '
+                          'in number of weights ({} vs {}).'.format(
+                              len(symbolic_weights), len(weight_values)))
+          continue
+        raise ValueError('Layer #' + str(k) + ' (named "' + layer.name +
+                         '") expects ' + str(len(symbolic_weights)) +
+                         ' weight(s), but the saved weights' + ' have ' +
+                         str(len(weight_values)) + ' element(s).')
+      # Set values.
+      for i in range(len(weight_values)):
+        if backend.int_shape(symbolic_weights[i]) != weight_values[i].shape:
+          if skip_mismatch:
+            logging.warning('Skipping loading of weights for '
+                            'layer {}'.format(layer.name) + ' due to '
+                            'mismatch in shape ({} vs {}).'.format(
+                                symbolic_weights[i].shape,
+                                weight_values[i].shape))
+            continue
+          raise ValueError('Layer #' + str(k) +' (named "' + layer.name +
+                           '"), weight ' + str(symbolic_weights[i]) +
+                           ' has shape {}'.format(backend.int_shape(
+                               symbolic_weights[i])) +
+                           ', but the saved weight has shape ' +
+                           str(weight_values[i].shape) + '.')
+
+        else:
+          weight_value_tuples.append((symbolic_weights[i], weight_values[i]))
+  backend.batch_set_value(weight_value_tuples)
+
+
+def save_attributes_to_hdf5_group(group, name, data):
+  """Saves attributes (data) of the specified name into the HDF5 group.
+
+  This method deals with an inherent problem of HDF5 file which is not
+  able to store data larger than HDF5_OBJECT_HEADER_LIMIT bytes.
+
+  Args:
+      group: A pointer to a HDF5 group.
+      name: A name of the attributes to save.
+      data: Attributes data to store.
+
+  Raises:
+    RuntimeError: If any single attribute is too large to be saved.
+  """
+  # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
+  # because in that case even chunking the array would not make the saving
+  # possible.
+  bad_attributes = [x for x in data if len(x) > HDF5_OBJECT_HEADER_LIMIT]
+
+  # Expecting this to never be true.
+  if bad_attributes:
+    raise RuntimeError('The following attributes cannot be saved to HDF5 '
+                       'file because they are larger than %d bytes: %s' %
+                       (HDF5_OBJECT_HEADER_LIMIT, ', '.join(bad_attributes)))
+
+  data_npy = np.asarray(data)
+
+  num_chunks = 1
+  chunked_data = np.array_split(data_npy, num_chunks)
+
+  # This will never loop forever thanks to the test above.
+  while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data):
+    num_chunks += 1
+    chunked_data = np.array_split(data_npy, num_chunks)
+
+  if num_chunks > 1:
+    for chunk_id, chunk_data in enumerate(chunked_data):
+      group.attrs['%s%d' % (name, chunk_id)] = chunk_data
+  else:
+    group.attrs[name] = data
+
+
+def load_attributes_from_hdf5_group(group, name):
+  """Loads attributes of the specified name from the HDF5 group.
+
+  This method deals with an inherent problem
+  of HDF5 file which is not able to store
+  data larger than HDF5_OBJECT_HEADER_LIMIT bytes.
+
+  Args:
+      group: A pointer to a HDF5 group.
+      name: A name of the attributes to load.
+
+  Returns:
+      data: Attributes data.
+  """
+  if name in group.attrs:
+    data = [
+        n.decode('utf8') if hasattr(n, 'decode') else n
+        for n in group.attrs[name]
+    ]
+  else:
+    data = []
+    chunk_id = 0
+    while '%s%d' % (name, chunk_id) in group.attrs:
+      data.extend([
+          n.decode('utf8') if hasattr(n, 'decode') else n
+          for n in group.attrs['%s%d' % (name, chunk_id)]
+      ])
+      chunk_id += 1
+  return data
+
+
+def _legacy_weights(layer):
+  """DO NOT USE.
+
+  For legacy reason, the layer.weights was in the order of
+  [self.trainable_weights + self.non_trainable_weights], and this order was
+  used for preserving the weights in h5 format. The new order of layer.weights
+  are the same as layer.get_weights() which is more intuitive for user. To
+  keep supporting the existing saved h5 file, this method should be used to
+  save/load weights. In future version, we will delete this method and
+  introduce a breaking change for h5 and stay with the new order for weights.
+
+  Args:
+    layer: a `tf.keras.Model` or `tf.keras.layers.Layer` instance.
+
+  Returns:
+    A list of variables with the order of trainable_weights, followed by
+      non_trainable_weights.
+  """
+  weights = layer.trainable_weights + layer.non_trainable_weights
+  if any(not isinstance(w, variables_module.Variable) for w in weights):
+    raise NotImplementedError(
+        'Save or restore weights that is not an instance of `tf.Variable` is '
+        'not supported in h5, use `save_format=\'tf\'` instead. Got a model '
+        'or layer {} with weights {}'.format(layer.__class__.__name__, weights))
+  return weights

tensorflow/python/keras/saving/model_config.py

@@ -0,0 +1,100 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# pylint: disable=protected-access
+"""Functions that save the model's config into different formats."""
+
+from tensorflow.python.keras.saving.saved_model import json_utils
+
+
+def model_from_config(config, custom_objects=None):
+  """Instantiates a Keras model from its config.
+
+  Usage:
+  ```
+  # for a Functional API model
+  tf.keras.Model().from_config(model.get_config())
+
+  # for a Sequential model
+  tf.keras.Sequential().from_config(model.get_config())
+  ```
+
+  Args:
+      config: Configuration dictionary.
+      custom_objects: Optional dictionary mapping names
+          (strings) to custom classes or functions to be
+          considered during deserialization.
+
+  Returns:
+      A Keras model instance (uncompiled).
+
+  Raises:
+      TypeError: if `config` is not a dictionary.
+  """
+  if isinstance(config, list):
+    raise TypeError('`model_from_config` expects a dictionary, not a list. '
+                    'Maybe you meant to use '
+                    '`Sequential.from_config(config)`?')
+  from tensorflow.python.keras.layers import deserialize  # pylint: disable=g-import-not-at-top
+  return deserialize(config, custom_objects=custom_objects)
+
+
+def model_from_yaml(yaml_string, custom_objects=None):
+  """Parses a yaml model configuration file and returns a model instance.
+
+  Note: Since TF 2.6, this method is no longer supported and will raise a
+  RuntimeError.
+
+  Args:
+      yaml_string: YAML string or open file encoding a model configuration.
+      custom_objects: Optional dictionary mapping names
+          (strings) to custom classes or functions to be
+          considered during deserialization.
+
+  Returns:
+      A Keras model instance (uncompiled).
+
+  Raises:
+      RuntimeError: announces that the method poses a security risk
+  """
+  raise RuntimeError(
+      'Method `model_from_yaml()` has been removed due to security risk of '
+      'arbitrary code execution. Please use `Model.to_json()` and '
+      '`model_from_json()` instead.'
+  )
+
+
+def model_from_json(json_string, custom_objects=None):
+  """Parses a JSON model configuration string and returns a model instance.
+
+  Usage:
+
+  >>> model = tf.keras.Sequential([
+  ...     tf.keras.layers.Dense(5, input_shape=(3,)),
+  ...     tf.keras.layers.Softmax()])
+  >>> config = model.to_json()
+  >>> loaded_model = tf.keras.models.model_from_json(config)
+
+  Args:
+      json_string: JSON string encoding a model configuration.
+      custom_objects: Optional dictionary mapping names
+          (strings) to custom classes or functions to be
+          considered during deserialization.
+
+  Returns:
+      A Keras model instance (uncompiled).
+  """
+  config = json_utils.decode(json_string)
+  from tensorflow.python.keras.layers import deserialize  # pylint: disable=g-import-not-at-top
+  return deserialize(config, custom_objects=custom_objects)

tensorflow/python/keras/saving/save.py

@@ -0,0 +1,209 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Keras model saving code."""
+import os
+from tensorflow.python import tf2
+from tensorflow.python.keras.saving import hdf5_format
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import load as saved_model_load
+from tensorflow.python.keras.saving.saved_model import load_context
+from tensorflow.python.keras.saving.saved_model import save as saved_model_save
+from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.keras.utils.io_utils import path_to_string
+
+# pylint: disable=g-import-not-at-top
+try:
+  # Disable loading HDF5 plugins from a default path and prevent ZDI-CAN-25480.
+  # Importing h5py prior to importing tensorflow will restore the old behavior.
+  os.environ.setdefault('HDF5_PLUGIN_PATH', os.devnull)
+  import h5py
+except ImportError:
+  h5py = None
+# pylint: enable=g-import-not-at-top
+
+
+def save_model(model,
+               filepath,
+               overwrite=True,
+               include_optimizer=True,
+               save_format=None,
+               signatures=None,
+               options=None,
+               save_traces=True):
+  # pylint: disable=line-too-long
+  """Saves a model as a TensorFlow SavedModel or HDF5 file.
+
+  See the [Serialization and Saving guide](https://keras.io/guides/serialization_and_saving/)
+  for details.
+
+  Usage:
+
+  >>> model = tf.keras.Sequential([
+  ...     tf.keras.layers.Dense(5, input_shape=(3,)),
+  ...     tf.keras.layers.Softmax()])
+  >>> model.save('/tmp/model')
+  >>> loaded_model = tf.keras.models.load_model('/tmp/model')
+  >>> x = tf.random.uniform((10, 3))
+  >>> assert np.allclose(model.predict(x), loaded_model.predict(x))
+
+  The SavedModel and HDF5 file contains:
+
+  - the model's configuration (topology)
+  - the model's weights
+  - the model's optimizer's state (if any)
+
+  Thus models can be reinstantiated in the exact same state, without any of the
+  code used for model definition or training.
+
+  Note that the model weights may have different scoped names after being
+  loaded. Scoped names include the model/layer names, such as
+  `"dense_1/kernel:0"`. It is recommended that you use the layer properties to
+  access specific variables, e.g. `model.get_layer("dense_1").kernel`.
+
+  __SavedModel serialization format__
+
+  Keras SavedModel uses `tf.saved_model.save` to save the model and all
+  trackable objects attached to the model (e.g. layers and variables). The model
+  config, weights, and optimizer are saved in the SavedModel. Additionally, for
+  every Keras layer attached to the model, the SavedModel stores:
+
+    * the config and metadata -- e.g. name, dtype, trainable status
+    * traced call and loss functions, which are stored as TensorFlow subgraphs.
+
+  The traced functions allow the SavedModel format to save and load custom
+  layers without the original class definition.
+
+  You can choose to not save the traced functions by disabling the `save_traces`
+  option. This will decrease the time it takes to save the model and the
+  amount of disk space occupied by the output SavedModel. If you enable this
+  option, then you _must_ provide all custom class definitions when loading
+  the model. See the `custom_objects` argument in `tf.keras.models.load_model`.
+
+  Args:
+      model: Keras model instance to be saved.
+      filepath: One of the following:
+        - String or `pathlib.Path` object, path where to save the model
+        - `h5py.File` object where to save the model
+      overwrite: Whether we should overwrite any existing model at the target
+        location, or instead ask the user with a manual prompt.
+      include_optimizer: If True, save optimizer's state together.
+      save_format: Either 'tf' or 'h5', indicating whether to save the model
+        to Tensorflow SavedModel or HDF5. Defaults to 'tf' in TF 2.X, and 'h5'
+        in TF 1.X.
+      signatures: Signatures to save with the SavedModel. Applicable to the 'tf'
+        format only. Please see the `signatures` argument in
+        `tf.saved_model.save` for details.
+      options: (only applies to SavedModel format) `tf.saved_model.SaveOptions`
+        object that specifies options for saving to SavedModel.
+      save_traces: (only applies to SavedModel format) When enabled, the
+        SavedModel will store the function traces for each layer. This
+        can be disabled, so that only the configs of each layer are stored.
+        Defaults to `True`. Disabling this will decrease serialization time and
+        reduce file size, but it requires that all custom layers/models
+        implement a `get_config()` method.
+
+  Raises:
+      ImportError: If save format is hdf5, and h5py is not available.
+  """
+  # pylint: enable=line-too-long
+  from tensorflow.python.keras.engine import sequential  # pylint: disable=g-import-not-at-top
+
+  default_format = 'tf' if tf2.enabled() else 'h5'
+  save_format = save_format or default_format
+
+  filepath = path_to_string(filepath)
+
+  # If the user has not already called fit or built the underlying metrics, we
+  # should do that before saving to ensure the metric names have all
+  # appropriate name transformations applied.
+  saving_utils.try_build_compiled_arguments(model)
+
+  if (save_format == 'h5' or
+      (h5py is not None and isinstance(filepath, h5py.File)) or
+      saving_utils.is_hdf5_filepath(filepath)):
+    # TODO(b/130258301): add utility method for detecting model type.
+    if (not model._is_graph_network and  # pylint:disable=protected-access
+        not isinstance(model, sequential.Sequential)):
+      raise NotImplementedError(
+          'Saving the model to HDF5 format requires the model to be a '
+          'Functional model or a Sequential model. It does not work for '
+          'subclassed models, because such models are defined via the body of '
+          'a Python method, which isn\'t safely serializable. Consider saving '
+          'to the Tensorflow SavedModel format (by setting save_format="tf") '
+          'or using `save_weights`.')
+    hdf5_format.save_model_to_hdf5(
+        model, filepath, overwrite, include_optimizer)
+  else:
+    with generic_utils.SharedObjectSavingScope():
+      saved_model_save.save(model, filepath, overwrite, include_optimizer,
+                            signatures, options, save_traces)
+
+
+def load_model(filepath, custom_objects=None, compile=True, options=None):  # pylint: disable=redefined-builtin
+  """Loads a model saved via `model.save()`.
+
+  Usage:
+
+  >>> model = tf.keras.Sequential([
+  ...     tf.keras.layers.Dense(5, input_shape=(3,)),
+  ...     tf.keras.layers.Softmax()])
+  >>> model.save('/tmp/model')
+  >>> loaded_model = tf.keras.models.load_model('/tmp/model')
+  >>> x = tf.random.uniform((10, 3))
+  >>> assert np.allclose(model.predict(x), loaded_model.predict(x))
+
+  Note that the model weights may have different scoped names after being
+  loaded. Scoped names include the model/layer names, such as
+  `"dense_1/kernel:0"`. It is recommended that you use the layer properties to
+  access specific variables, e.g. `model.get_layer("dense_1").kernel`.
+
+  Args:
+      filepath: One of the following:
+          - String or `pathlib.Path` object, path to the saved model
+          - `h5py.File` object from which to load the model
+      custom_objects: Optional dictionary mapping names
+          (strings) to custom classes or functions to be
+          considered during deserialization.
+      compile: Boolean, whether to compile the model
+          after loading.
+      options: Optional `tf.saved_model.LoadOptions` object that specifies
+        options for loading from SavedModel.
+
+  Returns:
+      A Keras model instance. If the original model was compiled, and saved with
+      the optimizer, then the returned model will be compiled. Otherwise, the
+      model will be left uncompiled. In the case that an uncompiled model is
+      returned, a warning is displayed if the `compile` argument is set to
+      `True`.
+
+  Raises:
+      ImportError: if loading from an hdf5 file and h5py is not available.
+      IOError: In case of an invalid savefile.
+  """
+  with generic_utils.SharedObjectLoadingScope():
+    with generic_utils.CustomObjectScope(custom_objects or {}):
+      with load_context.load_context(options):
+        if (h5py is not None and
+            (isinstance(filepath, h5py.File) or h5py.is_hdf5(filepath))):
+          return hdf5_format.load_model_from_hdf5(filepath, custom_objects,
+                                                  compile)
+
+        filepath = path_to_string(filepath)
+        if isinstance(filepath, str):
+          return saved_model_load.load(filepath, compile, options)
+
+  raise IOError(
+      'Unable to load model. Filepath is not an hdf5 file (or h5py is not '
+      'available) or SavedModel.')

tensorflow/python/keras/saving/saved_model/BUILD

@@ -0,0 +1,70 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# Description:
+#   Keras saving and loading files for SavedModel.
+
+load("@xla//third_party/rules_python/python:defs.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/python/keras/saving:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "load_context",
+    srcs = ["load_context.py"],
+    strict_deps = False,
+    visibility = ["//tensorflow/python/distribute:__pkg__"],
+)
+
+py_library(
+    name = "saved_model",
+    srcs = [
+        "base_serialization.py",
+        "constants.py",
+        "json_utils.py",
+        "layer_serialization.py",
+        "load.py",
+        "load_context.py",
+        "metric_serialization.py",
+        "model_serialization.py",
+        "network_serialization.py",
+        "save.py",
+        "save_impl.py",
+        "serialized_attributes.py",
+        "utils.py",
+    ],
+    srcs_version = "PY3",
+    strict_deps = False,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:signature_def_utils",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//tensorflow/python/saved_model:utils",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_decorator_py",
+    ],
+)

tensorflow/python/keras/saving/saved_model/README.md

@@ -0,0 +1,131 @@
+# Keras SavedModel
+
+For questions, feedback, and feature requests please file a bug/contact kathywu@
+
+## TensorFlow Core SavedModel implementation
+
+In TensorFlow 2.0, all saving and loading implementations revolve around the
+object graph generated from a root trackable object, and all trackable objects
+connected to it through attributes. Program building blocks such as variables,
+assets, and tables, and high level objects like Optimizers and Layers all
+subclass the trackable class. Other objects like TensorFlow functions and
+concrete functions are also saved as nodes in the object graph. When loading a
+SavedModel, the object graph is used to recreate the structure of the original
+object.
+
+Please see the links below for more details:
+
+- [Saved Model Guide](https://www.tensorflow.org/guide/saved_model)
+- [Checkpoint Guide](https://www.tensorflow.org/guide/checkpoint)
+
+## Keras SavedModel implementation
+
+### Overview
+
+Keras object serialization is built on top of the core serialization.
+
+All attributes that impact model execution or inspection are saved to the
+SavedModel to allow the model to be recreated. These attributes are divided into
+three categories:
+
+1. python properties (e.g., layer name, layer config)
+2. objects (e.g. data structures like list of variables or layers)
+3. functions (e.g. call function, loss functions)
+
+Trackable objects and TensorFlow functions are represented as nodes in the
+trackable object graph, and each node in the graph stores information about
+their python properties.
+
+Since many attributes in Keras Layers/Models are not Trackable objects or
+tf.functions, these attributes are wrapped as trackable objects/tf.functions at
+serialization time. For example, `layer.variables` is implemented as a python
+property that appends the lists of trainable/nontrainable variables. During
+serialization, a new Trackable List object is created and saved to the
+`variables` attribute. Another example is the call function. Most models do not
+decorate their call function with `tf.function`, since Keras will take care of
+the graph/function management. When the model is saved, the call function is
+wrapped in a `tf.function` and added to the `__call__` attribute.
+
+
+### `keras_api` attribute
+
+Many attributes are only relevant for revivability. Instead of attaching these
+directly to the exported object, they are saved to a new `keras_api` trackable
+object that is then attached to the exported object. This avoids cluttering the
+exported object with objects/functions that are only used by the Keras library.
+
+For example, `__call__` and `call_and_return_conditional_losses` are functions
+saved for all models. The `__call__` function is attached directly to the
+exported object, while `call_and_return_conditional_losses` is attached to a
+separate object. Say a user saves the model, then loads the SavedModel using the
+core loader (tf.saved_model.load which does not rely on the Keras library to
+revive the model).
+
+The loaded object will have a structure that looks like:
+
+```
+  loaded object -- __call__
+                -- keras_api -- __call__
+                             -- call_and_return_conditional_losses
+```
+
+The two call functions may be accessed through:
+
+  - `loaded.__call__` or `loaded.keras_api.__call__`
+  - `loaded.keras_api.call_and_return_conditional_losses`.
+
+
+### Saving details
+
+Keras Layers use a helper abstract class and an attribute validator class to
+define and standardize the serialization implementation:
+
+- [`SerializationImpl`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/base_serialization.py):
+Ensures that layer python properties are saved as a serialized JSON string in
+the metadata field, and gathers all attributes to save with the Keras object.
+Please see the docstrings in each of the abstract methods/properties to see what
+is required.
+- [`SerializedAttributes`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/serialized_attributes.py?):
+Tracks all of the attributes that must be saved with a Keras object. Objects and
+functions may be specified to be "keras_only", meaning that they will only
+appear in the `keras_api` attribute.
+
+The base `Layer` serialization is defined in
+[`layer_serialization.py`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/layer_serialization.py).
+See `LayerAttributes` and `LayerSerializationImpl`.
+
+**Adding a new attribute to base Layer SavedModel**
+
+1. Add a new attributes to `LayerAttributes`.
+2. Modify `LayerSerializationImpl` internal methods:
+
+   a. If adding a python property, add the key-value item to the dictionary
+   returned by `_python_properties_internal`
+
+   b.If adding a new object/function, modify the dictionary returned by
+   `_get_serialized_attributes_internal`.
+
+
+**Adding custom serialization for a Layer subclass.**
+
+1. Create a new attribute validator by copying `LayerAttributes`, and add any
+new attributes to serialize.
+2. Subclass `LayerSerializationImpl`
+3. Implement `_python_properties_internal` and/or
+`_get_serialized_attributes_internal` to return the new attributes.
+
+Unless you are modifying the loader (see section below on loading), please keep
+the `object_identifier` the same.
+
+These instructions also carry over for modifying
+[Model](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/model_serialization.py)
+and
+[Network](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/network_serialization.py)
+serialization.
+
+
+### Loading details
+
+TODO(kathywu): Will write this section when the loading code is moved into
+\*_serialization.py files.
+

tensorflow/python/keras/saving/saved_model/base_serialization.py

@@ -0,0 +1,134 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Helper classes that list&validate all attributes to serialize to SavedModel."""
+
+import abc
+
+from tensorflow.python.keras.saving.saved_model import json_utils
+from tensorflow.python.keras.saving.saved_model import utils
+
+
+class SavedModelSaver(object, metaclass=abc.ABCMeta):
+  """Saver defining the methods and properties used to serialize Keras objects.
+  """
+
+  def __init__(self, obj):
+    self.obj = obj
+
+  @abc.abstractproperty
+  def object_identifier(self):
+    """String stored in object identifier field in the SavedModel proto.
+
+    Returns:
+      A string with the object identifier, which is used at load time.
+    """
+    raise NotImplementedError
+
+  @property
+  def tracking_metadata(self):
+    """String stored in metadata field in the SavedModel proto.
+
+    Returns:
+      A serialized JSON storing information necessary for recreating this layer.
+    """
+    # TODO(kathywu): check that serialized JSON can be loaded (e.g., if an
+    # object is in the python property)
+    return json_utils.Encoder().encode(self.python_properties)
+
+  def trackable_children(self, serialization_cache):
+    """Lists all Trackable children connected to this object."""
+    if not utils.should_save_traces():
+      return {}
+
+    children = self.objects_to_serialize(serialization_cache)
+    children.update(self.functions_to_serialize(serialization_cache))
+    return children
+
+  @abc.abstractproperty
+  def python_properties(self):
+    """Returns dictionary of python properties to save in the metadata.
+
+    This dictionary must be serializable and deserializable to/from JSON.
+
+    When loading, the items in this dict are used to initialize the object and
+    define attributes in the revived object.
+    """
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def objects_to_serialize(self, serialization_cache):
+    """Returns dictionary of extra checkpointable objects to serialize.
+
+    See `functions_to_serialize` for an explanation of this function's
+    effects.
+
+    Args:
+      serialization_cache: Dictionary passed to all objects in the same object
+        graph during serialization.
+
+    Returns:
+        A dictionary mapping attribute names to checkpointable objects.
+    """
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def functions_to_serialize(self, serialization_cache):
+    """Returns extra functions to include when serializing a Keras object.
+
+    Normally, when calling exporting an object to SavedModel, only the
+    functions and objects defined by the user are saved. For example:
+
+    ```
+    obj = tf.Module()
+    obj.v = tf.Variable(1.)
+
+    @tf.function
+    def foo(...): ...
+
+    obj.foo = foo
+
+    w = tf.Variable(1.)
+
+    tf.saved_model.save(obj, 'path/to/saved/model')
+    loaded = tf.saved_model.load('path/to/saved/model')
+
+    loaded.v  # Variable with the same value as obj.v
+    loaded.foo  # Equivalent to obj.foo
+    loaded.w  # AttributeError
+    ```
+
+    Assigning trackable objects to attributes creates a graph, which is used for
+    both checkpointing and SavedModel serialization.
+
+    When the graph generated from attribute tracking is insufficient, extra
+    objects and functions may be added at serialization time. For example,
+    most models do not have their call function wrapped with a @tf.function
+    decorator. This results in `model.call` not being saved. Since Keras objects
+    should be revivable from the SavedModel format, the call function is added
+    as an extra function to serialize.
+
+    This function and `objects_to_serialize` is called multiple times when
+    exporting to SavedModel. Please use the cache to avoid generating new
+    functions and objects. A fresh cache is created for each SavedModel export.
+
+    Args:
+      serialization_cache: Dictionary passed to all objects in the same object
+        graph during serialization.
+
+    Returns:
+        A dictionary mapping attribute names to `Function` or
+        `ConcreteFunction`.
+    """
+    raise NotImplementedError

tensorflow/python/keras/saving/saved_model/constants.py

@@ -0,0 +1,47 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Constants for Keras SavedModel serialization."""
+
+# Namespace used to store all attributes added during serialization.
+# e.g. the list of layers can be accessed using `loaded.keras_api.layers`, in an
+# object loaded from `tf.saved_model.load()`.
+KERAS_ATTR = 'keras_api'
+
+# Keys for the serialization cache.
+# Maps to the keras serialization dict {Layer --> SerializedAttributes object}
+KERAS_CACHE_KEY = 'keras_serialized_attributes'
+
+
+# Name of Keras metadata file stored in the SavedModel.
+SAVED_METADATA_PATH = 'keras_metadata.pb'
+
+# Names of SavedObject Keras identifiers.
+INPUT_LAYER_IDENTIFIER = '_tf_keras_input_layer'
+LAYER_IDENTIFIER = '_tf_keras_layer'
+METRIC_IDENTIFIER = '_tf_keras_metric'
+MODEL_IDENTIFIER = '_tf_keras_model'
+NETWORK_IDENTIFIER = '_tf_keras_network'
+RNN_LAYER_IDENTIFIER = '_tf_keras_rnn_layer'
+SEQUENTIAL_IDENTIFIER = '_tf_keras_sequential'
+
+KERAS_OBJECT_IDENTIFIERS = (
+    INPUT_LAYER_IDENTIFIER,
+    LAYER_IDENTIFIER,
+    METRIC_IDENTIFIER,
+    MODEL_IDENTIFIER,
+    NETWORK_IDENTIFIER,
+    RNN_LAYER_IDENTIFIER,
+    SEQUENTIAL_IDENTIFIER,
+)

tensorflow/python/keras/saving/saved_model/json_utils.py

@@ -0,0 +1,144 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils for creating and loading the Layer metadata for SavedModel.
+
+These are required to retain the original format of the build input shape, since
+layers and models may have different build behaviors depending on if the shape
+is a list, tuple, or TensorShape. For example, Network.build() will create
+separate inputs if the given input_shape is a list, and will create a single
+input if the given shape is a tuple.
+"""
+
+import collections
+import enum
+import json
+import numpy as np
+import wrapt
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import type_spec_registry
+from tensorflow.python.types import internal
+
+
+class Encoder(json.JSONEncoder):
+  """JSON encoder and decoder that handles TensorShapes and tuples."""
+
+  def default(self, obj):  # pylint: disable=method-hidden
+    """Encodes objects for types that aren't handled by the default encoder."""
+    if isinstance(obj, tensor_shape.TensorShape):
+      items = obj.as_list() if obj.rank is not None else None
+      return {'class_name': 'TensorShape', 'items': items}
+    return get_json_type(obj)
+
+  def encode(self, obj):
+    return super(Encoder, self).encode(_encode_tuple(obj))
+
+
+def _encode_tuple(x):
+  if isinstance(x, tuple):
+    return {'class_name': '__tuple__',
+            'items': tuple(_encode_tuple(i) for i in x)}
+  elif isinstance(x, list):
+    return [_encode_tuple(i) for i in x]
+  elif isinstance(x, dict):
+    return {key: _encode_tuple(value) for key, value in x.items()}
+  else:
+    return x
+
+
+def decode(json_string):
+  return json.loads(json_string, object_hook=_decode_helper)
+
+
+def _decode_helper(obj):
+  """A decoding helper that is TF-object aware."""
+  if isinstance(obj, dict) and 'class_name' in obj:
+    if obj['class_name'] == 'TensorShape':
+      return tensor_shape.TensorShape(obj['items'])
+    elif obj['class_name'] == 'TypeSpec':
+      return type_spec_registry.lookup(obj['type_spec'])._deserialize(  # pylint: disable=protected-access
+          _decode_helper(obj['serialized']))
+    elif obj['class_name'] == '__tuple__':
+      return tuple(_decode_helper(i) for i in obj['items'])
+    elif obj['class_name'] == '__ellipsis__':
+      return Ellipsis
+  return obj
+
+
+def get_json_type(obj):
+  """Serializes any object to a JSON-serializable structure.
+
+  Args:
+      obj: the object to serialize
+
+  Returns:
+      JSON-serializable structure representing `obj`.
+
+  Raises:
+      TypeError: if `obj` cannot be serialized.
+  """
+  # if obj is a serializable Keras class instance
+  # e.g. optimizer, layer
+  if hasattr(obj, 'get_config'):
+    return {'class_name': obj.__class__.__name__, 'config': obj.get_config()}
+
+  # if obj is any numpy type
+  if type(obj).__module__ == np.__name__:
+    if isinstance(obj, np.ndarray):
+      return obj.tolist()
+    else:
+      return obj.item()
+
+  # misc functions (e.g. loss function)
+  if callable(obj):
+    return obj.__name__
+
+  # if obj is a python 'type'
+  if type(obj).__name__ == type.__name__:
+    return obj.__name__
+
+  if isinstance(obj, tensor_shape.Dimension):
+    return obj.value
+
+  if isinstance(obj, tensor_shape.TensorShape):
+    return obj.as_list()
+
+  if isinstance(obj, dtypes.DType):
+    return obj.name
+
+  if isinstance(obj, collections.abc.Mapping):
+    return dict(obj)
+
+  if obj is Ellipsis:
+    return {'class_name': '__ellipsis__'}
+
+  if isinstance(obj, wrapt.ObjectProxy):
+    return obj.__wrapped__
+
+  if isinstance(obj, internal.TypeSpec):
+    try:
+      type_spec_name = type_spec_registry.get_name(type(obj))
+      return {'class_name': 'TypeSpec', 'type_spec': type_spec_name,
+              'serialized': obj._serialize()}  # pylint: disable=protected-access
+    except ValueError:
+      raise ValueError('Unable to serialize {} to JSON, because the TypeSpec '
+                       'class {} has not been registered.'
+                       .format(obj, type(obj)))
+
+  if isinstance(obj, enum.Enum):
+    return obj.value
+
+  raise TypeError('Not JSON Serializable:', obj)

tensorflow/python/keras/saving/saved_model/layer_serialization.py

@@ -0,0 +1,175 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Classes and functions implementing Layer SavedModel serialization."""
+
+from tensorflow.python.keras.mixed_precision import policy
+from tensorflow.python.keras.saving.saved_model import base_serialization
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import save_impl
+from tensorflow.python.keras.saving.saved_model import serialized_attributes
+from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.util import nest
+
+
+class LayerSavedModelSaver(base_serialization.SavedModelSaver):
+  """Implements Layer SavedModel serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.LAYER_IDENTIFIER
+
+  @property
+  def python_properties(self):
+    # TODO(kathywu): Add python property validator
+    return self._python_properties_internal()
+
+  def _python_properties_internal(self):
+    """Returns dictionary of all python properties."""
+    # TODO(kathywu): Add support for metrics serialization.
+    # TODO(kathywu): Synchronize with the keras spec (go/keras-json-spec) once
+    # the python config serialization has caught up.
+    metadata = dict(
+        name=self.obj.name,
+        trainable=self.obj.trainable,
+        expects_training_arg=self.obj._expects_training_arg,  # pylint: disable=protected-access
+        dtype=policy.serialize(self.obj._dtype_policy),  # pylint: disable=protected-access
+        batch_input_shape=getattr(self.obj, '_batch_input_shape', None),
+        stateful=self.obj.stateful,
+        must_restore_from_config=self.obj._must_restore_from_config,  # pylint: disable=protected-access
+    )
+
+    metadata.update(get_serialized(self.obj))
+    if self.obj.input_spec is not None:
+      # Layer's input_spec has already been type-checked in the property setter.
+      metadata['input_spec'] = nest.map_structure(
+          lambda x: generic_utils.serialize_keras_object(x) if x else None,
+          self.obj.input_spec)
+    if (self.obj.activity_regularizer is not None and
+        hasattr(self.obj.activity_regularizer, 'get_config')):
+      metadata['activity_regularizer'] = generic_utils.serialize_keras_object(
+          self.obj.activity_regularizer)
+    if self.obj._build_input_shape is not None:  # pylint: disable=protected-access
+      metadata['build_input_shape'] = self.obj._build_input_shape  # pylint: disable=protected-access
+    return metadata
+
+  def objects_to_serialize(self, serialization_cache):
+    return (self._get_serialized_attributes(
+        serialization_cache).objects_to_serialize)
+
+  def functions_to_serialize(self, serialization_cache):
+    return (self._get_serialized_attributes(
+        serialization_cache).functions_to_serialize)
+
+  def _get_serialized_attributes(self, serialization_cache):
+    """Generates or retrieves serialized attributes from cache."""
+    keras_cache = serialization_cache.setdefault(constants.KERAS_CACHE_KEY, {})
+    if self.obj in keras_cache:
+      return keras_cache[self.obj]
+
+    serialized_attr = keras_cache[self.obj] = (
+        serialized_attributes.SerializedAttributes.new(self.obj))
+
+    if (save_impl.should_skip_serialization(self.obj) or
+        self.obj._must_restore_from_config):  # pylint: disable=protected-access
+      return serialized_attr
+
+    object_dict, function_dict = self._get_serialized_attributes_internal(
+        serialization_cache)
+
+    serialized_attr.set_and_validate_objects(object_dict)
+    serialized_attr.set_and_validate_functions(function_dict)
+    return serialized_attr
+
+  def _get_serialized_attributes_internal(self, serialization_cache):
+    """Returns dictionary of serialized attributes."""
+    objects = save_impl.wrap_layer_objects(self.obj, serialization_cache)
+    functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
+    # Attribute validator requires that the default save signature is added to
+    # function dict, even if the value is None.
+    functions['_default_save_signature'] = None
+    return objects, functions
+
+
+# TODO(kathywu): Move serialization utils (and related utils from
+# generic_utils.py) to a separate file.
+def get_serialized(obj):
+  with generic_utils.skip_failed_serialization():
+    # Store the config dictionary, which may be used when reviving the object.
+    # When loading, the program will attempt to revive the object from config,
+    # and if that fails, the object will be revived from the SavedModel.
+    return generic_utils.serialize_keras_object(obj)
+
+
+class InputLayerSavedModelSaver(base_serialization.SavedModelSaver):
+  """InputLayer serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.INPUT_LAYER_IDENTIFIER
+
+  @property
+  def python_properties(self):
+
+    return dict(
+        class_name=type(self.obj).__name__,
+        name=self.obj.name,
+        dtype=self.obj.dtype,
+        sparse=self.obj.sparse,
+        ragged=self.obj.ragged,
+        batch_input_shape=self.obj._batch_input_shape,  # pylint: disable=protected-access
+        config=self.obj.get_config())
+
+  def objects_to_serialize(self, serialization_cache):
+    return {}
+
+  def functions_to_serialize(self, serialization_cache):
+    return {}
+
+
+class RNNSavedModelSaver(LayerSavedModelSaver):
+  """RNN layer serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.RNN_LAYER_IDENTIFIER
+
+  def _get_serialized_attributes_internal(self, serialization_cache):
+    objects, functions = (
+        super(RNNSavedModelSaver, self)._get_serialized_attributes_internal(
+            serialization_cache))
+    states = data_structures.wrap_or_unwrap(self.obj.states)
+    # SaveModel require all the objects to be Trackable when saving.
+    # If the states is still a tuple after wrap_or_unwrap, it means it doesn't
+    # contain any trackable item within it, eg empty tuple or (None, None) for
+    # stateless ConvLSTM2D. We convert them to list so that wrap_or_unwrap can
+    # make it a Trackable again for saving. When loaded, ConvLSTM2D is
+    # able to handle the tuple/list conversion.
+    if isinstance(states, tuple):
+      states = data_structures.wrap_or_unwrap(list(states))
+    objects['states'] = states
+    return objects, functions
+
+
+class IndexLookupLayerSavedModelSaver(LayerSavedModelSaver):
+  """Index lookup layer serialization."""
+
+  @property
+  def python_properties(self):
+    # TODO(kathywu): Add python property validator
+    metadata = self._python_properties_internal()
+    if metadata['config'].get('has_static_table', False):
+      metadata['config']['vocabulary'] = None
+    return metadata

tensorflow/python/keras/saving/saved_model/load_context.py

@@ -0,0 +1,63 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Context for storing options for loading a SavedModel."""
+
+import contextlib
+import threading
+
+
+class LoadContext(threading.local):
+  """A context for loading a model."""
+
+  def __init__(self):
+    super(LoadContext, self).__init__()
+    self._entered_load_context = []
+    self._load_options = None
+
+  def set_load_options(self, load_options):
+    self._load_options = load_options
+    self._entered_load_context.append(True)
+
+  def clear_load_options(self):
+    self._load_options = None
+    self._entered_load_context.pop()
+
+  def load_options(self):
+    return self._load_options
+
+  def in_load_context(self):
+    return self._entered_load_context
+
+
+_load_context = LoadContext()
+
+
+@contextlib.contextmanager
+def load_context(load_options):
+  _load_context.set_load_options(load_options)
+  try:
+    yield
+  finally:
+    _load_context.clear_load_options()
+
+
+def get_load_options():
+  """Returns the load options under a load context."""
+  return _load_context.load_options()
+
+
+def in_load_context():
+  """Returns whether under a load context."""
+  return _load_context.in_load_context()

tensorflow/python/keras/saving/saved_model/metric_serialization.py

@@ -0,0 +1,43 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Classes and functions implementing Metrics SavedModel serialization."""
+
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import layer_serialization
+from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.trackable import data_structures
+
+
+class MetricSavedModelSaver(layer_serialization.LayerSavedModelSaver):
+  """Metric serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.METRIC_IDENTIFIER
+
+  def _python_properties_internal(self):
+    metadata = dict(
+        class_name=generic_utils.get_registered_name(type(self.obj)),
+        name=self.obj.name,
+        dtype=self.obj.dtype)
+    metadata.update(layer_serialization.get_serialized(self.obj))
+    if self.obj._build_input_shape is not None:  # pylint: disable=protected-access
+      metadata['build_input_shape'] = self.obj._build_input_shape  # pylint: disable=protected-access
+    return metadata
+
+  def _get_serialized_attributes_internal(self, unused_serialization_cache):
+    return (dict(variables=data_structures.wrap_or_unwrap(self.obj.variables)),
+            dict())  # TODO(b/135550038): save functions to enable saving
+                     # custom metrics.

tensorflow/python/keras/saving/saved_model/model_serialization.py

@@ -0,0 +1,63 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Classes and functions implementing to Model SavedModel serialization."""
+
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import layer_serialization
+from tensorflow.python.keras.saving.saved_model import save_impl
+
+
+class ModelSavedModelSaver(layer_serialization.LayerSavedModelSaver):
+  """Model SavedModel serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.MODEL_IDENTIFIER
+
+  def _python_properties_internal(self):
+    metadata = super(ModelSavedModelSaver, self)._python_properties_internal()
+    # Network stateful property is dependent on the child layers.
+    metadata.pop('stateful')
+    metadata['is_graph_network'] = self.obj._is_graph_network  # pylint: disable=protected-access
+    metadata['save_spec'] = self.obj._get_save_spec(dynamic_batch=False)  # pylint: disable=protected-access
+
+    metadata.update(
+        saving_utils.model_metadata(
+            self.obj, include_optimizer=True, require_config=False))
+    return metadata
+
+  def _get_serialized_attributes_internal(self, serialization_cache):
+    default_signature = None
+
+    # Create a default signature function if this is the only object in the
+    # cache (i.e. this is the root level object).
+    if len(serialization_cache[constants.KERAS_CACHE_KEY]) == 1:
+      default_signature = save_impl.default_save_signature(self.obj)
+
+    # Other than the default signature function, all other attributes match with
+    # the ones serialized by Layer.
+    objects, functions = (
+        super(ModelSavedModelSaver, self)._get_serialized_attributes_internal(
+            serialization_cache))
+    functions['_default_save_signature'] = default_signature
+    return objects, functions
+
+
+class SequentialSavedModelSaver(ModelSavedModelSaver):
+
+  @property
+  def object_identifier(self):
+    return constants.SEQUENTIAL_IDENTIFIER

tensorflow/python/keras/saving/saved_model/network_serialization.py

@@ -0,0 +1,27 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Classes and functions implementing to Network SavedModel serialization."""
+
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import model_serialization
+
+
+# FunctionalModel serialization is pretty much the same as Model serialization.
+class NetworkSavedModelSaver(model_serialization.ModelSavedModelSaver):
+  """Network serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.NETWORK_IDENTIFIER

tensorflow/python/keras/saving/saved_model/save.py

@@ -0,0 +1,124 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Keras SavedModel serialization."""
+
+import os
+
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.protobuf import saved_metadata_pb2
+from tensorflow.python.keras.protobuf import versions_pb2
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import save_impl
+from tensorflow.python.keras.saving.saved_model import utils
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.keras.utils.io_utils import ask_to_proceed_with_overwrite
+from tensorflow.python.platform import gfile
+from tensorflow.python.saved_model import save as save_lib
+
+# To avoid circular dependencies between keras/engine and keras/saving,
+# code in keras/saving must delay imports.
+
+base_layer = LazyLoader(
+    "base_layer", globals(),
+    "tensorflow.python.keras.engine.base_layer")
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+
+
+def save(model, filepath, overwrite, include_optimizer, signatures=None,
+         options=None, save_traces=True):
+  """Saves a model as a SavedModel to the filepath.
+
+  Args:
+    model: Keras model instance to be saved.
+    filepath: String path to save the model.
+    overwrite: whether to overwrite the existing filepath.
+    include_optimizer: If True, save the model's optimizer state.
+    signatures: Signatures to save with the SavedModel. Applicable to the 'tf'
+      format only. Please see the `signatures` argument in `tf.saved_model.save`
+      for details.
+    options: (only applies to SavedModel format) `tf.saved_model.SaveOptions`
+      object that specifies options for saving to SavedModel.
+    save_traces: (only applies to SavedModel format) When enabled, the
+      SavedModel will store the function traces for each layer. This
+      can be disabled, so that only the configs of each layer are stored.
+      Defaults to `True`. Disabling this will decrease serialization time
+      and reduce file size, but it requires that all custom layers/models
+      implement a `get_config()` method.
+
+  Raises:
+    ValueError: if the model's inputs have not been defined.
+  """
+  # If file exists and should not be overwritten.
+  if not overwrite and os.path.exists(filepath):
+    proceed = ask_to_proceed_with_overwrite(filepath)
+    if not proceed:
+      return
+
+  if save_traces:
+    if save_impl.should_skip_serialization(model):
+      saving_utils.raise_model_input_error(model)
+
+  if not include_optimizer:
+    orig_optimizer = model.optimizer
+    model.optimizer = None
+    # TODO(b/180760306) Change to del model.optimizer if Layer's __delattr__
+    # calls AutoTrackable's __delattr__.
+    model._delete_tracking("optimizer")  # pylint: disable=protected-access
+
+  # Trace all functions and signatures with `training=0` instead of using an
+  # already-set learning phase placeholder.
+  # This is needed for compatibility reasons until learning phase setting
+  # is removed from the public apis.
+  with K.deprecated_internal_learning_phase_scope(0):
+    with utils.keras_option_scope(save_traces):
+      saved_nodes, node_paths = save_lib.save_and_return_nodes(
+          model, filepath, signatures, options)
+
+    # Save all metadata to a separate file in the SavedModel directory.
+    metadata = generate_keras_metadata(saved_nodes, node_paths)
+
+  with gfile.GFile(
+      os.path.join(filepath, constants.SAVED_METADATA_PATH), "wb") as w:
+    w.write(metadata.SerializeToString(deterministic=True))
+
+  if not include_optimizer:
+    model.optimizer = orig_optimizer
+
+
+def generate_keras_metadata(saved_nodes, node_paths):
+  """Constructs a KerasMetadata proto with the metadata of each keras object."""
+  metadata = saved_metadata_pb2.SavedMetadata()
+
+  for node_id, node in enumerate(saved_nodes):
+    if isinstance(node, base_layer.Layer):
+      path = node_paths[node]
+      if not path:
+        node_path = "root"
+      else:
+        node_path = "root.{}".format(
+            ".".join([ref.name for ref in path]))
+
+      metadata.nodes.add(
+          node_id=node_id,
+          node_path=node_path,
+          version=versions_pb2.VersionDef(
+              producer=1, min_consumer=1, bad_consumers=[]),
+          identifier=node._object_identifier,  # pylint: disable=protected-access
+          metadata=node._tracking_metadata)  # pylint: disable=protected-access
+
+  return metadata

tensorflow/python/keras/saving/saved_model/save_impl.py

@@ -0,0 +1,734 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Keras SavedModel serialization.
+
+TODO (kathywu): Move to layer_serialization.py. Some model-specific logic should
+go to model_serialization.py.
+"""
+
+import functools
+import threading
+import weakref
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.engine import base_layer_utils
+from tensorflow.python.keras.engine import input_spec
+from tensorflow.python.keras.mixed_precision import autocast_variable
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import load as keras_load
+from tensorflow.python.keras.saving.saved_model import serialized_attributes
+from tensorflow.python.keras.saving.saved_model import utils
+from tensorflow.python.keras.utils import tf_contextlib
+from tensorflow.python.keras.utils import tf_inspect
+from tensorflow.python.keras.utils import tf_utils
+from tensorflow.python.keras.utils import version_utils
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_decorator
+
+
+# To avoid circular dependencies between keras/engine and keras/saving,
+# code in keras/saving must delay imports.
+
+# TODO(b/134426265): Switch back to single-quotes to match the rest of the file
+# once the issue with copybara is fixed.
+# pylint:disable=g-inconsistent-quotes
+base_layer = LazyLoader(
+    "base_layer", globals(),
+    "tensorflow.python.keras.engine.base_layer")
+metrics = LazyLoader("metrics", globals(),
+                     "tensorflow.python.keras.metrics")
+input_layer = LazyLoader(
+    "input_layer", globals(),
+    "tensorflow.python.keras.engine.input_layer")
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+sequential_lib = LazyLoader(
+    "sequential_lib", globals(),
+    "tensorflow.python.keras.engine.sequential")
+# pylint:enable=g-inconsistent-quotes
+
+
+def should_skip_serialization(layer):
+  """Skip serializing extra objects and functions if layer inputs aren't set."""
+  saved_model_input_spec_set = (isinstance(layer, training_lib.Model) and
+                                layer._saved_model_inputs_spec is not None)  # pylint: disable=protected-access
+  if not layer.built and not saved_model_input_spec_set:
+    logging.warning('Skipping full serialization of Keras layer {}, because '
+                    'it is not built.'.format(layer))
+    return True
+  return False
+
+
+def wrap_layer_objects(layer, serialization_cache):
+  """Returns extra trackable objects to attach to the serialized layer.
+
+  Args:
+    layer: Keras Layer object.
+    serialization_cache: Dictionary shared between all objects during
+      serialization.
+
+  Returns:
+    A dictionary containing all checkpointable objects from a
+    SerializedAttributes object. See LayerAttributes and ModelAttributes for
+    entire list of objects
+  """
+  # Wrap all regularization losses as tf.functions.
+  # First, generate list of all regularization losses in this layer and
+  # sublayers.
+  all_losses = layer._callable_losses[:]  # pylint: disable=protected-access
+  for child_layer in utils.list_all_layers(layer):
+    all_losses.extend(child_layer._callable_losses)  # pylint: disable=protected-access
+  # Next, wrap all loss functions as tf.functions. Use the serialization cache
+  # to store already-wrapped functions.
+  keras_loss_cache = serialization_cache.setdefault('keras_losses', {})
+  wrapped_loss_functions = []
+  for loss_fn in all_losses:
+    if loss_fn in keras_loss_cache:
+      wrapped_loss_functions.append(keras_loss_cache[loss_fn])
+    else:
+      wrapped_loss = _wrap_unconditional_loss(loss_fn, len(keras_loss_cache))
+      keras_loss_cache[loss_fn] = wrapped_loss
+      wrapped_loss_functions.append(wrapped_loss)
+  wrapped_layer_losses = [keras_loss_cache[fn]
+                          for fn in layer._callable_losses[:]]  # pylint: disable=protected-access
+
+  layer_metrics = data_structures.wrap_or_unwrap(
+      {m.name: m for m in layer._metrics})  # pylint: disable=protected-access
+  return dict(
+      variables=data_structures.wrap_or_unwrap(layer.variables),
+      trainable_variables=data_structures.wrap_or_unwrap(
+          layer.trainable_variables),
+      non_trainable_variables=data_structures.wrap_or_unwrap(
+          layer.non_trainable_variables),
+      layers=data_structures.wrap_or_unwrap(utils.list_all_layers(layer)),
+      metrics=data_structures.wrap_or_unwrap(layer.metrics),
+      regularization_losses=data_structures.wrap_or_unwrap(
+          wrapped_loss_functions),
+      layer_regularization_losses=data_structures.wrap_or_unwrap(
+          wrapped_layer_losses),
+      layer_metrics=layer_metrics)
+  # pylint: disable=protected-access
+
+
+def wrap_layer_functions(layer, serialization_cache):
+  """Returns dict of wrapped layer call function and losses in tf.functions.
+
+  Args:
+    layer: Keras Layer object.
+    serialization_cache: Dictionary shared between all objects during
+      serialization.
+
+  Returns:
+    A dictionary containing all keras tf.functions to serialize. See
+    LayerAttributes and ModelAttributes for the list of all attributes.
+  """
+  # Since Sequential models may be modified in place using model.add() or
+  # model.pop(), don't use saved functions.
+  if (isinstance(layer, keras_load.RevivedLayer) and
+      not isinstance(layer, sequential_lib.Sequential)):
+    return {fn_name: getattr(layer.keras_api, fn_name, None)
+            for fn_name in serialized_attributes.LayerAttributes.all_functions}
+
+  # Reset the losses of the layer and its children. The call function in each
+  # child layer is replaced with tf.functions.
+  original_fns = _replace_child_layer_functions(layer, serialization_cache)
+  original_losses = _reset_layer_losses(layer)
+
+  # Wrap all the layer call and activity regularizer functions.
+
+  # Use LayerCallCollection to ensure that all layer call functions (__call__,
+  # call with losses) are traced with the same inputs.
+  call_collection = LayerCallCollection(layer)
+  call_fn_with_losses = call_collection.add_function(
+      _wrap_call_and_conditional_losses(layer),
+      '{}_layer_call_and_return_conditional_losses'.format(layer.name),
+      # If any of this layer's child layers use the training arg, the traced
+      # call functions of this layer will have a training keyword argument. If
+      # the original layer does not expect the training arg, then it will have
+      # to be removed (by setting `match_layer_training_arg`).
+      match_layer_training_arg=True)
+  call_fn = call_collection.add_function(
+      _extract_outputs_from_fn(layer, call_fn_with_losses),
+      '{}_layer_call_fn'.format(layer.name),
+      # Since `call_fn` wraps call_fn_with_losses and not the original call
+      # function, `match_layer_training_arg` should be set to False.
+      match_layer_training_arg=False)
+
+  fns = {'call_and_return_conditional_losses': call_fn_with_losses,
+         '__call__': call_fn}
+
+  if layer._activity_regularizer is not None:  # pylint: disable=protected-access
+    fns['activity_regularizer_fn'] = _wrap_activity_regularizer(layer)
+    fns['call_and_return_all_conditional_losses'] = (
+        call_collection.add_function(
+            _append_activity_regularizer_loss(
+                layer, call_fn_with_losses, fns['activity_regularizer_fn']),
+            '{}_layer_call_and_return_all_conditional_losses'.format(
+                layer.name),
+            match_layer_training_arg=False))
+  else:
+    fns['activity_regularizer_fn'] = None
+    fns['call_and_return_all_conditional_losses'] = call_fn_with_losses
+
+  # Manually trigger traces before restoring the overwritten functions. The
+  # functions are traced within the layer call context to ensure that layer
+  # functions (e.g. add_loss) behave as though running in graph mode.
+  with tracing_scope():
+    call_collection.trace_with_input_signature()
+    with base_layer_utils.call_context().enter(
+        layer, inputs=None, build_graph=True, training=None, saving=True):
+      for fn in fns.values():
+        if fn is not None and fn.input_signature is not None:
+          if isinstance(fn, LayerCall):
+            fn = fn.wrapped_call
+          fn.get_concrete_function()
+
+  # Restore overwritten functions and losses
+  _restore_child_layer_functions(original_fns)
+  _restore_layer_losses(original_losses)
+
+  return fns
+
+
+def default_save_signature(layer):
+  original_losses = _reset_layer_losses(layer)
+  fn = saving_utils.trace_model_call(layer)
+  fn.get_concrete_function()
+  _restore_layer_losses(original_losses)
+  return fn
+
+
+def _replace_child_layer_functions(layer, serialization_cache):
+  """Replaces functions in the children layers with wrapped tf.functions.
+
+  This step allows functions from parent layers to reference the wrapped
+  functions from their children layers instead of retracing the ops.
+
+  This function also resets all losses stored in the layer. These are stored in
+  the returned dictionary. Use `_restore_child_layer_functions` to restore
+  the original attributes.
+
+  Args:
+    layer: Keras Layer object.
+    serialization_cache: Dictionary shared between all objects during
+      serialization.
+
+  Returns:
+    Dictionary mapping layer objects -> original functions and losses:
+      { Child layer 1: {
+          'losses': Original losses,
+          'call': Original call function
+          '_activity_regularizer': Original activity regularizer},
+        Child layer 2: ...
+      }
+  """
+  # pylint: disable=protected-access
+  original_fns = {}
+
+  def replace_layer_functions(child_layer, serialized_fns):
+    """Replaces layer call and activity regularizer with wrapped functions."""
+    original_fns[child_layer] = {
+        'call': child_layer.call,
+        '_activity_regularizer': child_layer._activity_regularizer
+    }
+    with utils.no_automatic_dependency_tracking_scope(child_layer):
+      try:
+        child_layer._activity_regularizer = serialized_fns.get(
+            'activity_regularizer_fn')
+      except AttributeError:
+        # Some layers have an unsettable activity regularizer.
+        pass
+      child_layer.call = utils.use_wrapped_call(
+          child_layer,
+          serialized_fns['call_and_return_conditional_losses'],
+          default_training_value=False)
+
+  def replace_metric_functions(child_layer, serialized_fns):
+    """Replaces metric functions with wrapped functions."""
+    original_fns[child_layer] = {
+        '__call__': child_layer.__call__,
+        'result': child_layer.result,
+        'update_state': child_layer.update_state
+    }
+    with utils.no_automatic_dependency_tracking_scope(child_layer):
+      child_layer.__call__ = serialized_fns['__call__']
+      child_layer.result = serialized_fns['result']
+      child_layer.update_state = serialized_fns['update_state']
+
+  for child_layer in utils.list_all_layers(layer):
+    if isinstance(child_layer, input_layer.InputLayer):
+      continue
+
+    if child_layer not in serialization_cache[constants.KERAS_CACHE_KEY]:
+      serialized_functions = (
+          child_layer._trackable_saved_model_saver._get_serialized_attributes(
+              serialization_cache).functions)
+    else:
+      serialized_functions = (
+          serialization_cache[constants.KERAS_CACHE_KEY][child_layer].functions)
+    if not serialized_functions:
+      # This indicates either:
+      #   - circular dependency, which means the current layer's functions
+      #     should be wrapped first.
+      #   - Child layer's inputs are not defined, so its functions have not been
+      #     wrapped. In this case, no replacement is necessary so move on to the
+      #     next child.
+      continue
+
+    if isinstance(child_layer, metrics.Metric):
+      replace_metric_functions(child_layer, serialized_functions)
+    else:
+      replace_layer_functions(child_layer, serialized_functions)
+
+  return original_fns
+  # pylint: enable=protected-access
+
+
+def _restore_child_layer_functions(original_fns):
+  """Restores attributes replaced with `_replace_child_layer_functions`."""
+  for child_layer, fns in original_fns.items():
+    with utils.no_automatic_dependency_tracking_scope(child_layer):
+      for fn_name, fn in fns.items():
+        try:
+          setattr(child_layer, fn_name, fn)  # pylint: disable=protected-access
+        except AttributeError:
+          pass  # In the case of _activity_regularizer, setting the attribute
+          # may be disallowed.
+
+
+# pylint: disable=protected-access
+def _reset_layer_losses(parent_layer):
+  """Resets losses of layer and its sublayers, and returns original losses."""
+  losses_dict = {}
+  for layer in utils.list_all_layers_and_sublayers(parent_layer):
+    losses_dict[layer] = {'losses': layer._losses[:],
+                          'eager_losses': layer._eager_losses[:]}
+    with utils.no_automatic_dependency_tracking_scope(layer):
+      layer._losses = []
+      layer._eager_losses = []
+  return losses_dict
+
+
+def _restore_layer_losses(losses_dict):
+  for layer in losses_dict:
+    with utils.no_automatic_dependency_tracking_scope(layer):
+      layer._losses = losses_dict[layer]['losses']
+      layer._eager_losses = losses_dict[layer]['eager_losses']
+# pylint: enable=protected-access
+
+
+class LayerTracingContext(threading.local):
+
+  def __init__(self):
+    super(LayerTracingContext, self).__init__()
+    self.enable_call_tracing = False
+    self.trace_queue = []
+
+_thread_local_data = LayerTracingContext()
+
+
+@tf_contextlib.contextmanager
+def tracing_scope():
+  """Enables tracing scope."""
+  # This enables the LayerCallCollection's tracing mechanism to trace all call
+  # functions in the collection.
+  previous_value = _thread_local_data.enable_call_tracing
+  previous_queue = _thread_local_data.trace_queue
+  try:
+    _thread_local_data.enable_call_tracing = True
+    _thread_local_data.trace_queue = []
+    yield
+  finally:
+    # Run traces from the queue.
+    while _thread_local_data.trace_queue:
+      fn, args, kwargs, training = _thread_local_data.trace_queue.pop()
+      if training is not None:
+        with K.deprecated_internal_learning_phase_scope(training):
+          fn.get_concrete_function(*args, **kwargs)
+      else:
+        fn.get_concrete_function(*args, **kwargs)
+    _thread_local_data.trace_queue = previous_queue
+    _thread_local_data.enable_call_tracing = previous_value
+
+
+def add_trace_to_queue(fn, args, kwargs, training=None):
+  if tracing_enabled():
+    _thread_local_data.trace_queue.append(
+        (fn, args[:], kwargs.copy(), training))
+
+
+def tracing_enabled():
+  """Whether to add extra traces to the queue."""
+  return _thread_local_data.enable_call_tracing
+
+
+class LayerCallCollection(object):
+  """Groups wrapped layer call functions.
+
+  This is used to ensure that all layer call functions are traced with the same
+  inputs-
+    - call
+    - call_and_return_conditional_losses
+    - call_and_return_all_conditional_losses
+  """
+
+  def __init__(self, layer):
+    self.layer = layer
+
+    self.layer_call_method = _get_layer_call_method(layer)
+    self._expects_training_arg = utils.layer_uses_training_bool(layer)
+    self._training_arg_index = utils.get_training_arg_index(
+        self.layer_call_method)
+
+    # If the layer call function has kwargs, then the traced function cannot
+    # have an input signature.
+    arg_spec = tf_inspect.getfullargspec(self.layer_call_method)
+    self._has_kwargs = bool(self._expects_training_arg or
+                            arg_spec.defaults or
+                            arg_spec.kwonlyargs or
+                            arg_spec.varkw)
+
+    self._input_signature = self._generate_input_signature(layer)
+    self._functions = weakref.WeakValueDictionary()
+
+    # Get the input argument name from the args.
+    args = arg_spec.args
+    if tf_inspect.ismethod(self.layer_call_method):
+      args = args[1:]
+    self._input_arg_name = args[0] if args else 'inputs'
+
+  def _generate_input_signature(self, layer):
+    """Inspects layer object and returns the inferred input signature.
+
+    Args:
+      layer: Layer object.
+
+    Returns:
+      List of possibly nested TensorSpecs of the layer call function inputs.
+      The list does not contain the `training` argument.
+    """
+    if (isinstance(layer.call, def_function.Function) and
+        layer.call.input_signature is not None):
+      return layer.call.input_signature
+    elif isinstance(layer, training_lib.Model):
+      return saving_utils.model_input_signature(layer)
+    elif (layer.input_spec is not None and
+          layer._use_input_spec_as_call_signature):  # pylint: disable=protected-access
+
+      def to_tensor_spec_or_none(x):
+        spec = input_spec.to_tensor_spec(x, layer._compute_dtype)  # pylint: disable=protected-access
+        # If the shape is too general (e.g. multiple dimensions are allowed),
+        # return None so that separate functions can be generated for each
+        # inferred input signature.
+        # TODO(b/134962016): currently partial signatures are not supported.
+        if spec.shape == tensor_shape.TensorShape(None):
+          return None
+        return spec
+      input_signature = [nest.map_structure(
+          to_tensor_spec_or_none, layer.input_spec)]
+
+      return input_signature
+    else:
+      return None
+
+  def add_trace(self, *args, **kwargs):
+    """Traces all functions with the same args and kwargs.
+
+    Args:
+      *args: Positional args passed to the original function.
+      **kwargs: Keyword args passed to the original function.
+    """
+    args = list(args)
+    kwargs = kwargs.copy()
+
+    for fn in self._functions.values():
+      # TODO(kathywu): Replace arguments with broader shapes defined in the
+      # input signature.
+      if self._expects_training_arg:
+        def trace_with_training(value, fn=fn):
+          utils.set_training_arg(value, self._training_arg_index, args, kwargs)
+          add_trace_to_queue(fn, args, kwargs, value)
+
+        trace_with_training(True)
+        trace_with_training(False)
+      else:
+        add_trace_to_queue(fn, args, kwargs)
+
+  @property
+  def fn_input_signature(self):
+    """Returns input signature for the wrapped layer call function."""
+    if self._has_kwargs:
+      # Input signatures may only describe tensor arguments and kwargs are not
+      # supported.
+      return None
+    if None in nest.flatten(self._input_signature):
+      # TODO(b/134962016): If input signature cannot be partially defined.
+      return None
+    return self._input_signature
+
+  def training_arg_was_passed(self, args, kwargs):
+    if not self.layer._expects_training_arg and self._expects_training_arg:  # pylint: disable=protected-access
+      return (utils.get_training_arg(self._training_arg_index, args, kwargs)
+              is not None)
+    else:
+      return self.layer._call_arg_was_passed(  # pylint: disable=protected-access
+          'training', args, kwargs, inputs_in_args=True)
+
+  def get_training_arg_value(self, args, kwargs):
+    if not self.layer._expects_training_arg and self._expects_training_arg:  # pylint: disable=protected-access
+      return utils.get_training_arg(self._training_arg_index, args, kwargs)
+    else:
+      return self.layer._get_call_arg_value(  # pylint: disable=protected-access
+          'training', args, kwargs, inputs_in_args=True)
+
+  def get_input_arg_value(self, args, kwargs):
+    return self.layer._get_call_arg_value(  # pylint: disable=protected-access
+        self._input_arg_name, args, kwargs, inputs_in_args=True)
+
+  def _maybe_wrap_with_training_arg(self, call_fn, match_layer_training_arg):
+    """Wraps call function with added training argument if necessary."""
+    if not self.layer._expects_training_arg and self._expects_training_arg:  # pylint: disable=protected-access
+      # Add training arg to wrapper function.
+      arg_spec = tf_inspect.getfullargspec(call_fn)
+      args = arg_spec.args + ['training']
+      defaults = list(arg_spec.defaults or [])
+      defaults.append(False)
+      new_arg_spec = tf_inspect.FullArgSpec(
+          args=args,
+          varargs=arg_spec.varargs,
+          varkw=arg_spec.varkw,
+          defaults=defaults,
+          kwonlyargs=arg_spec.kwonlyargs,
+          kwonlydefaults=arg_spec.kwonlydefaults,
+          annotations=arg_spec.annotations)
+
+      # Set new training arg index
+      self._training_arg_index = len(args) - 1
+      if tf_inspect.ismethod(call_fn):
+        self._training_arg_index -= 1
+
+      def wrap_with_training_arg(*args, **kwargs):
+        if match_layer_training_arg:
+          # Remove the training value, since the original call_fn does not
+          # expect a training arg. Instead, the training value will be
+          # propagated using the call context created in LayerCall.
+          args = list(args)
+          kwargs = kwargs.copy()
+          utils.remove_training_arg(self._training_arg_index, args, kwargs)
+        return call_fn(*args, **kwargs)
+
+      return tf_decorator.make_decorator(
+          target=call_fn,
+          decorator_func=wrap_with_training_arg,
+          decorator_argspec=new_arg_spec)
+
+    return call_fn
+
+  def add_function(self, call_fn, name, match_layer_training_arg):
+    """Adds a layer call function to the collection.
+
+    Args:
+      call_fn: a python function
+      name: Name of call function
+      match_layer_training_arg: If True, removes the `training` from the
+        function arguments when calling `call_fn`.
+
+    Returns:
+      LayerCall (tf.function)
+    """
+    fn = LayerCall(
+        self,
+        self._maybe_wrap_with_training_arg(call_fn, match_layer_training_arg),
+        name,
+        input_signature=self.fn_input_signature)
+    self._functions[name] = fn.wrapped_call
+    return fn
+
+  def trace_with_input_signature(self):
+    """Trace with the layer/models inferred input signature if possible."""
+    if (None not in nest.flatten(self._input_signature) and self._has_kwargs):
+      # Manually add traces for layers that have keyword arguments and have
+      # a fully defined input signature.
+      self.add_trace(*self._input_signature)
+
+
+def _filtered_inputs(inputs):
+  return list(filter(tf_utils.is_tensor_or_variable, nest.flatten(inputs)))
+
+
+def layer_call_wrapper(call_collection, method, name):
+  """Ensures layer losses are kept the same, and runs method in call context."""
+
+  # Create wrapper that deals with losses and call context.
+  def wrapper(*args, **kwargs):
+    """Calls method within call context."""
+    layer = call_collection.layer
+    training = None
+    inputs = _filtered_inputs([args, kwargs])
+    # pylint: disable=protected-access
+    if (args or kwargs) and call_collection.training_arg_was_passed(
+        args, kwargs):
+      training = call_collection.get_training_arg_value(args, kwargs)
+    # pylint: enable=protected-access
+    original_losses = _reset_layer_losses(layer)
+    with base_layer_utils.call_context().enter(
+        layer, inputs=inputs, build_graph=False, training=training,
+        saving=True):
+      with autocast_variable.enable_auto_cast_variables(
+          layer._compute_dtype_object):  # pylint: disable=protected-access
+        ret = method(*args, **kwargs)
+    _restore_layer_losses(original_losses)
+    return ret
+
+  # Rename to `name`, since tf.function doesn't have a name argument. Without
+  # this, all functions returned by this method will be named "call", which
+  # would be a nightmare to debug.
+  fn = tf_decorator.make_decorator(target=method, decorator_func=wrapper)
+  fn.__name__ = name
+  return fn
+
+
+class LayerCall(object):
+  """Function that triggers traces of other functions in the same collection."""
+
+  def __init__(self, call_collection, call_fn, name, input_signature):
+    """Initializes a LayerCall object.
+
+    Args:
+      call_collection: a LayerCallCollection, which contains the other layer
+        call functions (e.g. call_with_conditional_losses, call). These
+        functions should be traced with the same arguments.
+      call_fn: A call function.
+      name: Name of the call function.
+      input_signature: Input signature of call_fn (can be None).
+    """
+    self.call_collection = call_collection
+    self.input_signature = input_signature
+    self.wrapped_call = def_function.function(
+        layer_call_wrapper(call_collection, call_fn, name),
+        input_signature=input_signature)
+    self.original_layer_call = call_collection.layer_call_method
+
+  def _maybe_trace(self, args, kwargs):
+    # Trigger traces of other call functions + extra training-arg traces.
+    if tracing_enabled():
+      self.call_collection.add_trace(*args, **kwargs)
+
+  def __call__(self, *args, **kwargs):
+    self._maybe_trace(args, kwargs)
+    return self.wrapped_call(*args, **kwargs)
+
+  def get_concrete_function(self, *args, **kwargs):
+    self._maybe_trace(args, kwargs)
+    return self.wrapped_call.get_concrete_function(*args, **kwargs)
+
+
+def _wrap_call_and_conditional_losses(layer):
+  """Wraps call function that returns a tuple of (outputs, losses).
+
+  The losses returned are conditional on the inputs passed to the call function.
+  Unconditional losses (e.g. weight regularizeration) are wrapped separately.
+
+  Args:
+    layer: a Keras layer object
+
+  Returns:
+    python call function that returns outputs and conditional losses -- excludes
+    activity regularizer
+  """
+  # Create function that generates both outputs and losses
+  layer_call = _get_layer_call_method(layer)
+  def call_and_return_conditional_losses(*args, **kwargs):
+    """Returns layer (call_output, conditional losses) tuple."""
+    call_output = layer_call(*args, **kwargs)
+    if version_utils.is_v1_layer_or_model(layer):
+      conditional_losses = layer.get_losses_for(
+          _filtered_inputs([args, kwargs]))
+    else:
+      conditional_losses = [
+          l for l in layer.losses if not hasattr(l, '_unconditional_loss')
+      ]
+    return call_output, conditional_losses
+
+  return _create_call_fn_decorator(layer, call_and_return_conditional_losses)
+
+
+def _extract_outputs_from_fn(layer, call_and_return_conditional_losses):
+  """Returns a function that returns only call function outputs."""
+  if isinstance(layer, keras_load.RevivedLayer):
+    return layer.keras_api.__call__  # pylint: disable=protected-access
+  def call(inputs, *args, **kwargs):
+    return call_and_return_conditional_losses(inputs, *args, **kwargs)[0]
+  return _create_call_fn_decorator(layer, call)
+
+
+def _append_activity_regularizer_loss(
+    layer, call_fn_with_losses, activity_regularizer_fn):
+  """Appends activity regularizer loss to losses returned by the wrapped fn."""
+  def fn(inputs, *args, **kwargs):
+    outputs, losses = call_fn_with_losses(inputs, *args, **kwargs)
+    losses.append(activity_regularizer_fn(outputs))
+    return outputs, losses
+  return _create_call_fn_decorator(layer, fn)
+
+
+def _create_call_fn_decorator(layer, wrapped_call):
+  call_fn = _get_layer_call_method(layer)
+  fn, arg_spec = utils.maybe_add_training_arg(
+      call_fn, wrapped_call, layer._expects_training_arg,  # pylint: disable=protected-access
+      default_training_value=False)
+  return tf_decorator.make_decorator(
+      target=call_fn,
+      decorator_func=fn,
+      decorator_argspec=arg_spec)
+
+
+def _wrap_unconditional_loss(loss_fn, index):
+  """Wraps callable/unconditional loss, returning a serializable function."""
+  # Extract original loss function from partial function
+  fn = loss_fn.args[0] if isinstance(loss_fn, functools.partial) else loss_fn
+  if isinstance(fn, def_function.Function):
+    return fn
+  else:
+    return def_function.Function(
+        fn, 'loss_fn_{}'.format(index), input_signature=[])
+
+
+def _wrap_activity_regularizer(layer):
+  """Wraps the activity regularizer."""
+  # pylint: disable=protected-access
+  if isinstance(layer._activity_regularizer, def_function.Function):
+    return layer._activity_regularizer
+  return def_function.Function(
+      layer._activity_regularizer,
+      '{}_activity_regularizer'.format(layer.name),
+      input_signature=[
+          tensor_spec.TensorSpec(None, layer._compute_dtype or K.floatx())
+      ])
+  # pylint: enable=protected-access
+
+
+def _get_layer_call_method(layer):
+  if isinstance(layer.call, (def_function.Function)):
+    return layer.call.python_function
+  return layer.call

tensorflow/python/keras/saving/saved_model/serialized_attributes.py

@@ -0,0 +1,311 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Helper classes that list&validate all attributes to serialize to SavedModel.
+"""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import save_impl
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.trackable.autotrackable import AutoTrackable
+
+# TODO(b/134426265): Switch back to single-quotes to match the rest of the file
+# once the issue with copybara is fixed.
+# pylint:disable=g-inconsistent-quotes
+base_layer = LazyLoader(
+    "base_layer", globals(),
+    "tensorflow.python.keras.engine.base_layer")
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+metrics = LazyLoader("metrics", globals(),
+                     "tensorflow.python.keras.metrics")
+recurrent = LazyLoader(
+    "recurrent", globals(),
+    "tensorflow.python.keras.layers.recurrent")
+# pylint:enable=g-inconsistent-quotes
+
+
+class SerializedAttributes(object):
+  """Class that tracks and validates all serialization attributes.
+
+  Keras models contain many Python-defined components. For example, the
+  trainable_variable property lists the model's trainable variables by
+  recursively retrieving the trainable variables from each of the child layers.
+  Another example is model.call, a python function that calls child layers and
+  adds ops to the backend graph.
+
+  Only Tensorflow checkpointable objects and functions can be serialized to
+  SavedModel. Serializing a Keras model as-is results in a checkpointable object
+  that does not resemble a Keras model at all. Thus, extra checkpointable
+  objects and functions must be created during serialization.
+
+  **Defining new serialized attributes**
+  Child classes should be defined using:
+    SerializedAttributes.with_attributes(
+        'name', checkpointable_objects=[...], functions=[...], copy_from=[...])
+  This class is used to cache generated checkpointable objects and functions,
+  ensuring that new objects and functions are generated a single time.
+
+  **Usage during serialization**
+  Each Layer/Model object should have a corresponding instance of
+  SerializedAttributes. Create a new instance by calling
+  `SerializedAttributes.new(obj)`. Objects and functions may be saved using
+  `.set_and_validate_checkpointable_objects`/`.set_and_and_validate_functions`.
+  The properties `.checkpointable_objects` and `.functions` returns the cached
+  values.
+
+  **Adding/changing attributes to save to SavedModel**
+  1. Change the call to `SerializedAttributes.with_attributes` in the correct
+     class:
+     - CommonEndpoints: Base attributes to be added during serialization. If
+       these attributes are present in a Trackable object, it can be
+       deserialized to a Keras Model.
+     - LayerAttributes: Attributes to serialize for Layer objects.
+     - ModelAttributes: Attributes to serialize for Model objects.
+  2. Update class docstring
+  3. Update arguments to any calls to `set_and_validate_*`. For example, if
+     `call_raw_tensors` is added to the ModelAttributes function list, then
+     a `call_raw_tensors` function should be passed to
+     `set_and_validate_functions`.
+
+  **Common endpoints vs other attributes**
+  Only common endpoints are attached directly to the root object. Keras-specific
+  attributes are saved to a separate trackable object with the name "keras_api".
+  The number of objects attached to the root is limited because any naming
+  conflicts will cause user code to break.
+
+  Another reason is that this will only affect users who call
+  `tf.saved_model.load` instead of `tf.keras.models.load_model`. These are
+  advanced users who are likely to have defined their own tf.functions and
+  trackable objects. The added Keras-specific attributes are kept out of the way
+  in the "keras_api" namespace.
+
+  Properties defined in this class may be used to filter out keras-specific
+  attributes:
+  - `functions_to_serialize`: Returns dict of functions to attach to the root
+      object.
+  - `checkpointable_objects_to_serialize`: Returns dict of objects to attach to
+      the root object (including separate trackable object containing
+      keras-specific attributes)
+
+  All changes to the serialized attributes must be backwards-compatible, so
+  attributes should not be removed or modified without sufficient justification.
+  """
+
+  @staticmethod
+  def with_attributes(
+      name, checkpointable_objects=None, functions=None, copy_from=None):
+    """Creates a subclass with all attributes as specified in the arguments.
+
+    Args:
+      name: Name of subclass
+      checkpointable_objects: List of checkpointable objects to be serialized
+        in the SavedModel.
+      functions: List of functions to be serialized in the SavedModel.
+      copy_from: List of other SerializedAttributes subclasses. The returned
+        class will copy checkpoint objects/functions from each subclass.
+
+    Returns:
+      Child class with attributes as defined in the `checkpointable_objects`
+      and `functions` lists.
+    """
+    checkpointable_objects = checkpointable_objects or []
+    functions = functions or []
+
+    if copy_from is not None:
+      for cls in copy_from:
+        checkpointable_objects.extend(cls.all_checkpointable_objects)
+        functions.extend(cls.all_functions)
+
+    classdict = {
+        'all_checkpointable_objects': set(checkpointable_objects),
+        'all_functions': set(functions)}
+    return type(name, (SerializedAttributes,), classdict)
+
+  @staticmethod
+  def new(obj):
+    """Returns a new SerializedAttribute object."""
+    if isinstance(obj, training_lib.Model):
+      return ModelAttributes()
+    elif isinstance(obj, metrics.Metric):
+      return MetricAttributes()
+    elif isinstance(obj, recurrent.RNN):
+      return RNNAttributes()
+    elif isinstance(obj, base_layer.Layer):
+      return LayerAttributes()
+    else:
+      raise TypeError('Internal error during serialization: Expected Keras '
+                      'Layer object, got {} of type {}'.format(obj, type(obj)))
+
+  def __init__(self):
+    self._object_dict = {}
+    self._function_dict = {}
+    self._keras_trackable = AutoTrackable()
+
+  @property
+  def functions(self):
+    """Returns dictionary of all functions."""
+    return {key: value for key, value in self._function_dict.items()
+            if value is not None}
+
+  @property
+  def checkpointable_objects(self):
+    """Returns dictionary of all checkpointable objects."""
+    return {key: value for key, value in self._object_dict.items()
+            if value is not None}
+
+  @property
+  def functions_to_serialize(self):
+    """Returns functions to attach to the root object during serialization."""
+    functions = {}
+    for key, v in self.functions.items():
+      if key in CommonEndpoints.all_functions:
+        functions[key] = (v.wrapped_call if isinstance(v, save_impl.LayerCall)
+                          else v)
+    return functions
+
+  @property
+  def objects_to_serialize(self):
+    """Returns objects to attach to the root object during serialization."""
+    objects = {key: value for key, value in self.checkpointable_objects.items()
+               if key in CommonEndpoints.all_checkpointable_objects}
+    objects[constants.KERAS_ATTR] = self._keras_trackable
+    return objects
+
+  def set_and_validate_functions(self, function_dict):
+    """Saves function dictionary, and validates dictionary values."""
+    for key in self.all_functions:
+      if key in function_dict:
+        if (function_dict[key] is not None and  # Not all functions are required
+            not isinstance(function_dict[key],
+                           (def_function.Function, save_impl.LayerCall))):
+          raise ValueError(
+              'Function dictionary contained a non-function object: {} (for key'
+              ' {})'.format(function_dict[key], key))
+        fn = function_dict[key]
+        self._function_dict[key] = fn
+
+        # Extract TensorFlow `Function` from LayerCall.
+        tf_fn = fn.wrapped_call if isinstance(fn, save_impl.LayerCall) else fn
+        setattr(self._keras_trackable, key, tf_fn)
+      else:
+        raise ValueError('Function {} missing from serialized function dict.'
+                         .format(key))
+    return self.functions
+
+  def set_and_validate_objects(self, object_dict):
+    """Saves objects to a dictionary, and validates the values."""
+    for key in self.all_checkpointable_objects:
+      if key in object_dict:
+        if not isinstance(object_dict[key], trackable.Trackable):
+          raise ValueError(
+              'Object dictionary contained a non-trackable object: {} (for key'
+              ' {})'.format(object_dict[key], key))
+        self._object_dict[key] = object_dict[key]
+        setattr(self._keras_trackable, key, object_dict[key])
+      else:
+        raise ValueError(
+            'Object {} missing from serialized object dict.'.format(key))
+    return self.checkpointable_objects
+
+
+class CommonEndpoints(SerializedAttributes.with_attributes(
+    'CommonEndpoints',
+    checkpointable_objects=['variables', 'trainable_variables',
+                            'regularization_losses'],
+    functions=['__call__', 'call_and_return_all_conditional_losses',
+               '_default_save_signature'])):
+  """Common endpoints shared by all models loadable by Keras.
+
+  List of all attributes:
+    variables: List of all variables in the model and its sublayers.
+    trainable_variables: List of all trainable variables in the model and its
+      sublayers.
+    regularization_losses: List of all unconditional losses (losses not
+      dependent on the inputs) in the model and its sublayers.
+    __call__: Function that takes inputs and returns the outputs of the model
+      call function.
+    call_and_return_all_conditional_losses: Function that returns a tuple of
+      (call function outputs, list of all losses that depend on the inputs).
+    _default_save_signature: Traced model call function. This is only included
+      if the top level exported object is a Keras model.
+  """
+
+
+class LayerAttributes(SerializedAttributes.with_attributes(
+    'LayerAttributes',
+    checkpointable_objects=['non_trainable_variables', 'layers', 'metrics',
+                            'layer_regularization_losses', 'layer_metrics'],
+    functions=['call_and_return_conditional_losses', 'activity_regularizer_fn'],
+    copy_from=[CommonEndpoints]
+    )):
+  """Layer checkpointable objects + functions that are saved to the SavedModel.
+
+  List of all attributes:
+    All attributes from CommonEndpoints
+    non_trainable_variables: List of non-trainable variables in the layer and
+      its sublayers.
+    layers: List of all sublayers.
+    metrics: List of all metrics in the layer and its sublayers.
+    call_and_return_conditional_losses: Function that takes inputs and returns a
+      tuple of (outputs of the call function, list of input-dependent losses).
+      The list of losses excludes the activity regularizer function, which is
+      separate to allow the deserialized Layer object to define a different
+      activity regularizer.
+    activity_regularizer_fn: Callable that returns the activity regularizer loss
+    layer_regularization_losses: List of losses owned only by this layer.
+    layer_metrics: List of metrics owned by this layer.
+  """
+
+
+class ModelAttributes(SerializedAttributes.with_attributes(
+    'ModelAttributes',
+    copy_from=[LayerAttributes])):
+  """Model checkpointable objects + functions that are saved to the SavedModel.
+
+  List of all attributes:
+    All attributes from LayerAttributes (including CommonEndpoints)
+  """
+  # TODO(kathywu): Add attributes `compile_losses` and `compile_metrics`, which
+  #  list all losses and metrics defined by `model.compile`.
+
+
+class MetricAttributes(
+    SerializedAttributes.with_attributes(
+        'MetricAttributes',
+        checkpointable_objects=['variables'],
+        functions=[],
+    )):
+  """Attributes that are added to Metric objects when saved to SavedModel.
+
+  List of all attributes:
+    variables: list of all variables
+  """
+  pass
+
+
+class RNNAttributes(SerializedAttributes.with_attributes(
+    'RNNAttributes',
+    checkpointable_objects=['states'],
+    copy_from=[LayerAttributes])):
+  """RNN checkpointable objects + functions that are saved to the SavedModel.
+
+  List of all attributes:
+    All attributes from LayerAttributes (including CommonEndpoints)
+    states: List of state variables
+  """
+

tensorflow/python/keras/saving/saved_model/utils.py

@@ -0,0 +1,306 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions shared between SavedModel saving/loading implementations."""
+
+import itertools
+import threading
+import types
+
+from tensorflow.python.eager import context
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.engine import base_layer_utils
+from tensorflow.python.keras.utils import control_flow_util
+from tensorflow.python.keras.utils import tf_contextlib
+from tensorflow.python.keras.utils import tf_inspect
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.util import tf_decorator
+
+
+# pylint:disable=g-inconsistent-quotes
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+# pylint:enable=g-inconsistent-quotes
+
+
+def use_wrapped_call(layer, call_fn, default_training_value=None,
+                     return_method=False):
+  """Creates fn that adds the losses returned by call_fn & returns the outputs.
+
+  Args:
+    layer: A Keras layer object
+    call_fn: tf.function that takes layer inputs (and possibly a training arg),
+      and returns a tuple of (outputs, list of losses).
+    default_training_value: Default value of the training kwarg. If `None`, the
+      default is `K.learning_phase()`.
+    return_method: Whether to return a method bound to the layer.
+
+  Returns:
+    function that calls call_fn and returns the outputs. Losses returned by
+    call_fn are added to the layer losses.
+  """
+  expects_training_arg = layer_uses_training_bool(layer)
+  if hasattr(call_fn, 'original_layer_call'):  # call_fn is a LayerCall object
+    original_call = call_fn.original_layer_call
+    # In Python 3, callable objects are not compatible with inspect.getargspec
+    call_fn = call_fn.__call__
+  else:
+    original_call = call_fn
+  fn, arg_spec = maybe_add_training_arg(
+      original_call, call_fn, expects_training_arg, default_training_value)
+
+  def return_outputs_and_add_losses(*args, **kwargs):
+    """Returns the outputs from the layer call function, and adds the losses."""
+    if return_method:
+      args = args[1:]
+
+    outputs, losses = fn(*args, **kwargs)
+    layer.add_loss(losses, inputs=True)
+
+    # TODO(kathywu): This is a temporary hack. When a network of layers is
+    # revived from SavedModel, only the top-level layer will have losses. This
+    # causes issues in eager mode because the child layers may have graph losses
+    # (thus model.losses returns a mix of Eager and graph tensors). To fix this,
+    # whenever eager losses are added to one layer, add eager losses to all
+    # child layers. This causes `.losses` to only return eager losses.
+    # pylint: disable=protected-access
+    if context.executing_eagerly():
+      for i in layer._flatten_layers():
+        if i is not layer:
+          i._eager_losses = [base_layer_utils.REVIVED_LOSS_PLACEHOLDER]
+    # pylint: enable=protected-access
+    return outputs
+
+  decorated = tf_decorator.make_decorator(
+      target=call_fn,
+      decorator_func=return_outputs_and_add_losses,
+      decorator_argspec=arg_spec)
+
+  if return_method:
+    return types.MethodType(decorated, layer)
+  else:
+    return decorated
+
+
+def layer_uses_training_bool(layer):
+  """Returns whether this layer or any of its children uses the training arg."""
+  if layer._expects_training_arg:  # pylint: disable=protected-access
+    return True
+  visited = {layer}
+  to_visit = list_all_layers(layer)
+  while to_visit:
+    layer = to_visit.pop()
+    if layer in visited:
+      continue
+    if getattr(layer, '_expects_training_arg', True):
+      return True
+    visited.add(layer)
+    to_visit.extend(list_all_layers(layer))
+  return False
+
+
+def list_all_layers(obj):
+  if isinstance(obj, training_lib.Model):
+    # Handle special case of Sequential, which doesn't return
+    # the `Input` layer.
+    return obj.layers
+  else:
+    return list(obj._flatten_layers(include_self=False, recursive=False))  # pylint: disable=protected-access
+
+
+def list_all_layers_and_sublayers(obj):
+  s = set([obj])
+  s.update(itertools.chain.from_iterable(
+      list_all_layers_and_sublayers(layer) for layer in list_all_layers(obj)))
+  return s
+
+
+def maybe_add_training_arg(
+    original_call, wrapped_call, expects_training_arg, default_training_value):
+  """Decorate call and optionally adds training argument.
+
+  If a layer expects a training argument, this function ensures that 'training'
+  is present in the layer args or kwonly args, with the default training value.
+
+  Args:
+    original_call: Original call function.
+    wrapped_call: Wrapped call function.
+    expects_training_arg: Whether to include 'training' argument.
+    default_training_value: Default value of the training kwarg to include in
+      the arg spec. If `None`, the default is `K.learning_phase()`.
+
+  Returns:
+    Tuple of (
+      function that calls `wrapped_call` and sets the training arg,
+      Argspec of returned function or `None` if the argspec is unchanged)
+  """
+  if not expects_training_arg:
+    return wrapped_call, None
+  def wrap_with_training_arg(*args, **kwargs):
+    """Wrap the `wrapped_call` function, and set training argument."""
+    training_arg_index = get_training_arg_index(original_call)
+    training = get_training_arg(training_arg_index, args, kwargs)
+    if training is None:
+      training = default_training_value or K.learning_phase()
+
+    args = list(args)
+    kwargs = kwargs.copy()
+
+    def replace_training_and_call(training):
+      set_training_arg(training, training_arg_index, args, kwargs)
+      return wrapped_call(*args, **kwargs)
+
+    return control_flow_util.smart_cond(
+        training, lambda: replace_training_and_call(True),
+        lambda: replace_training_and_call(False))
+
+  # Create arg spec for decorated function. If 'training' is not defined in the
+  # args of the original arg spec, then add it to kwonlyargs.
+  arg_spec = tf_inspect.getfullargspec(original_call)
+  defaults = list(arg_spec.defaults) if arg_spec.defaults is not None else []
+
+  kwonlyargs = arg_spec.kwonlyargs
+  kwonlydefaults = arg_spec.kwonlydefaults or {}
+  # Add training arg if it does not exist, or set the default training value.
+  if 'training' not in arg_spec.args:
+    kwonlyargs.append('training')
+    kwonlydefaults['training'] = default_training_value
+  else:
+    index = arg_spec.args.index('training')
+    training_default_index = len(arg_spec.args) - index
+    if (arg_spec.defaults and
+        len(arg_spec.defaults) >= training_default_index and
+        defaults[-training_default_index] is None):
+      defaults[-training_default_index] = default_training_value
+
+  decorator_argspec = tf_inspect.FullArgSpec(
+      args=arg_spec.args,
+      varargs=arg_spec.varargs,
+      varkw=arg_spec.varkw,
+      defaults=defaults,
+      kwonlyargs=kwonlyargs,
+      kwonlydefaults=kwonlydefaults,
+      annotations=arg_spec.annotations)
+  return wrap_with_training_arg, decorator_argspec
+
+
+def get_training_arg_index(call_fn):
+  """Returns the index of 'training' in the layer call function arguments.
+
+  Args:
+    call_fn: Call function.
+
+  Returns:
+    - n: index of 'training' in the call function arguments.
+    - -1: if 'training' is not found in the arguments, but layer.call accepts
+          variable keyword arguments
+    - None: if layer doesn't expect a training argument.
+  """
+  argspec = tf_inspect.getfullargspec(call_fn)
+  if argspec.varargs:
+    # When there are variable args, training must be a keyword arg.
+    if 'training' in argspec.kwonlyargs or argspec.varkw:
+      return -1
+    return None
+  else:
+    # Try to find 'training' in the list of args or kwargs.
+    arg_list = argspec.args
+    if tf_inspect.ismethod(call_fn):
+      arg_list = arg_list[1:]
+
+    if 'training' in arg_list:
+      return arg_list.index('training')
+    elif 'training' in argspec.kwonlyargs or argspec.varkw:
+      return -1
+    return None
+
+
+def set_training_arg(training, index, args, kwargs):
+  if index is None or index < 0 or len(args) <= index:  # index is invalid
+    kwargs['training'] = training
+  else:
+    args[index] = training
+  return args, kwargs
+
+
+def get_training_arg(index, args, kwargs):
+  if index is None or index < 0 or len(args) <= index:  # index is invalid
+    return kwargs.get('training', None)
+  else:
+    return args[index]
+
+
+def remove_training_arg(index, args, kwargs):
+  if index is None or index < 0 or len(args) <= index:  # index is invalid
+    kwargs.pop('training', None)
+  else:
+    args.pop(index)
+
+
+class SaveOptionsContext(threading.local):
+
+  def __init__(self):
+    super(SaveOptionsContext, self).__init__()
+    self.save_traces = True
+
+
+_save_options_context = SaveOptionsContext()
+
+
+@tf_contextlib.contextmanager
+def keras_option_scope(save_traces):
+  previous_value = _save_options_context.save_traces
+  try:
+    _save_options_context.save_traces = save_traces
+    yield
+  finally:
+    _save_options_context.save_traces = previous_value
+
+
+def should_save_traces():
+  """Whether to trace layer functions-can be disabled in the save_traces arg."""
+  return _save_options_context.save_traces
+
+
+@tf_contextlib.contextmanager
+def no_automatic_dependency_tracking_scope(obj):
+  """A context that disables automatic dependency tracking when assigning attrs.
+
+  Objects that inherit from Autotrackable automatically creates dependencies
+  to trackable objects through attribute assignments, and wraps data structures
+  (lists or dicts) with trackable classes. This scope may be used to temporarily
+  disable this behavior. This works similar to the decorator
+  `no_automatic_dependency_tracking`.
+
+  Example usage:
+  ```
+  model = tf.keras.Model()
+  model.arr1 = []  # Creates a ListWrapper object
+  with no_automatic_dependency_tracking_scope(model):
+    model.arr2 = []  # Creates a regular, untracked python list
+  ```
+
+  Args:
+    obj: A trackable object.
+
+  Yields:
+    a scope in which the object doesn't track dependencies.
+  """
+  previous_value = getattr(obj, '_setattr_tracking', True)
+  obj._setattr_tracking = False  # pylint: disable=protected-access
+  try:
+    yield
+  finally:
+    obj._setattr_tracking = previous_value  # pylint: disable=protected-access

tensorflow/python/keras/saving/saved_model_experimental.py

@@ -0,0 +1,464 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Deprecated experimental Keras SavedModel implementation."""
+
+import os
+import warnings
+from tensorflow.python.checkpoint import graph_view
+from tensorflow.python.client import session
+from tensorflow.python.framework import ops
+from tensorflow.python.keras import backend
+from tensorflow.python.keras import optimizer_v1
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.keras.saving import model_config
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving import utils_v1 as model_utils
+from tensorflow.python.keras.utils import mode_keys
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+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 save as save_lib
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+
+# To avoid circular dependencies between keras/engine and keras/saving,
+# code in keras/saving must delay imports.
+
+# TODO(b/134426265): Switch back to single-quotes to match the rest of the file
+# once the issue with copybara is fixed.
+# pylint:disable=g-inconsistent-quotes
+metrics_lib = LazyLoader("metrics_lib", globals(),
+                         "tensorflow.python.keras.metrics")
+models_lib = LazyLoader("models_lib", globals(),
+                        "tensorflow.python.keras.models")
+sequential = LazyLoader(
+    "sequential", globals(),
+    "tensorflow.python.keras.engine.sequential")
+# pylint:enable=g-inconsistent-quotes
+
+
+# File name for json format of SavedModel.
+SAVED_MODEL_FILENAME_JSON = 'saved_model.json'
+
+
+def export_saved_model(model,
+                       saved_model_path,
+                       custom_objects=None,
+                       as_text=False,
+                       input_signature=None,
+                       serving_only=False):
+  """Exports a `tf.keras.Model` as a Tensorflow SavedModel.
+
+  Note that at this time, subclassed models can only be saved using
+  `serving_only=True`.
+
+  The exported `SavedModel` is a standalone serialization of Tensorflow objects,
+  and is supported by TF language APIs and the Tensorflow Serving system.
+  To load the model, use the function
+  `tf.keras.experimental.load_from_saved_model`.
+
+  The `SavedModel` contains:
+
+  1. a checkpoint containing the model weights.
+  2. a `SavedModel` proto containing the Tensorflow backend graph. Separate
+     graphs are saved for prediction (serving), train, and evaluation. If
+     the model has not been compiled, then only the graph computing predictions
+     will be exported.
+  3. the model's json config. If the model is subclassed, this will only be
+     included if the model's `get_config()` method is overwritten.
+
+  Example:
+
+  ```python
+  import tensorflow as tf
+
+  # Create a tf.keras model.
+  model = tf.keras.Sequential()
+  model.add(tf.keras.layers.Dense(1, input_shape=[10]))
+  model.summary()
+
+  # Save the tf.keras model in the SavedModel format.
+  path = '/tmp/simple_keras_model'
+  tf.keras.experimental.export_saved_model(model, path)
+
+  # Load the saved keras model back.
+  new_model = tf.keras.experimental.load_from_saved_model(path)
+  new_model.summary()
+  ```
+
+  Args:
+    model: A `tf.keras.Model` to be saved. If the model is subclassed, the flag
+      `serving_only` must be set to True.
+    saved_model_path: a string specifying the path to the SavedModel directory.
+    custom_objects: Optional dictionary mapping string names to custom classes
+      or functions (e.g. custom loss functions).
+    as_text: bool, `False` by default. Whether to write the `SavedModel` proto
+      in text format. Currently unavailable in serving-only mode.
+    input_signature: A possibly nested sequence of `tf.TensorSpec` objects, used
+      to specify the expected model inputs. See `tf.function` for more details.
+    serving_only: bool, `False` by default. When this is true, only the
+      prediction graph is saved.
+
+  Raises:
+    NotImplementedError: If the model is a subclassed model, and serving_only is
+      False.
+    ValueError: If the input signature cannot be inferred from the model.
+    AssertionError: If the SavedModel directory already exists and isn't empty.
+  """
+  warnings.warn('`tf.keras.experimental.export_saved_model` is deprecated'
+                'and will be removed in a future version. '
+                'Please use `model.save(..., save_format="tf")` or '
+                '`tf.keras.models.save_model(..., save_format="tf")`.')
+  if serving_only:
+    save_lib.save(
+        model,
+        saved_model_path,
+        signatures=saving_utils.trace_model_call(model, input_signature))
+  else:
+    _save_v1_format(model, saved_model_path, custom_objects, as_text,
+                    input_signature)
+
+  try:
+    _export_model_json(model, saved_model_path)
+  except NotImplementedError:
+    logging.warning('Skipped saving model JSON, subclassed model does not have '
+                    'get_config() defined.')
+
+
+def _export_model_json(model, saved_model_path):
+  """Saves model configuration as a json string under assets folder."""
+  model_json = model.to_json()
+  model_json_filepath = os.path.join(
+      _get_or_create_assets_dir(saved_model_path),
+      compat.as_text(SAVED_MODEL_FILENAME_JSON))
+  with gfile.Open(model_json_filepath, 'w') as f:
+    f.write(model_json)
+
+
+def _export_model_variables(model, saved_model_path):
+  """Saves model weights in checkpoint format under variables folder."""
+  _get_or_create_variables_dir(saved_model_path)
+  checkpoint_prefix = _get_variables_path(saved_model_path)
+  model.save_weights(checkpoint_prefix, save_format='tf', overwrite=True)
+  return checkpoint_prefix
+
+
+def _save_v1_format(model, path, custom_objects, as_text, input_signature):
+  """Exports model to v1 SavedModel format."""
+  if not model._is_graph_network:  # pylint: disable=protected-access
+    if isinstance(model, sequential.Sequential):
+      # If input shape is not directly set in the model, the exported model
+      # will infer the expected shapes of the input from the model.
+      if not model.built:
+        raise ValueError('Weights for sequential model have not yet been '
+                         'created. Weights are created when the Model is first '
+                         'called on inputs or `build()` is called with an '
+                         '`input_shape`, or the first layer in the model has '
+                         '`input_shape` during construction.')
+      # TODO(kathywu): Build the model with input_signature to create the
+      # weights before _export_model_variables().
+    else:
+      raise NotImplementedError(
+          'Subclassed models can only be exported for serving. Please set '
+          'argument serving_only=True.')
+
+  builder = saved_model_builder._SavedModelBuilder(path)  # pylint: disable=protected-access
+
+  # 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_variables(model, path)
+
+  # Export each mode.
+  # 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,
+                 'input_signature': input_signature}
+
+  has_saved_vars = False
+  if model.optimizer:
+    if isinstance(model.optimizer, (optimizer_v1.TFOptimizer,
+                                    optimizer_v2.OptimizerV2)):
+      _export_mode(mode_keys.ModeKeys.TRAIN, has_saved_vars, **export_args)
+      has_saved_vars = True
+      _export_mode(mode_keys.ModeKeys.TEST, 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(mode_keys.ModeKeys.PREDICT, has_saved_vars, **export_args)
+
+  builder.save(as_text)
+
+
+def _get_var_list(model):
+  """Returns list of all checkpointed saveable objects in the model."""
+  var_list, _, _ = graph_view.ObjectGraphView(model).serialize_object_graph()
+  return var_list
+
+
+def create_placeholder(spec):
+  return backend.placeholder(shape=spec.shape, dtype=spec.dtype, name=spec.name)
+
+
+def _export_mode(
+    mode, has_saved_vars, builder, model, custom_objects, checkpoint_path,
+    input_signature):
+  """Exports a model, and optionally saves new vars from the clone model.
+
+  Args:
+    mode: A `KerasModeKeys` 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.
+    input_signature: Nested TensorSpec containing the expected inputs. Can be
+      `None`, in which case the signature will be inferred from the model.
+
+  Raises:
+    ValueError: If the train/eval mode is being exported, but the model does
+      not have an optimizer.
+  """
+  compile_clone = (mode != mode_keys.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, backend.learning_phase_scope(
+      mode == mode_keys.ModeKeys.TRAIN):
+
+    if input_signature is None:
+      input_tensors = None
+    else:
+      input_tensors = nest.map_structure(create_placeholder, input_signature)
+
+    # Clone the model into blank graph. This will create placeholders for inputs
+    # and targets.
+    clone = models_lib.clone_and_build_model(
+        model, input_tensors=input_tensors, 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.compat.v1.train.get_global_step()`.
+    if compile_clone:
+      g.add_to_collection(ops.GraphKeys.GLOBAL_STEP, clone.optimizer.iterations)
+
+    # Extract update and train ops from train/test/predict functions.
+    train_op = None
+    if mode == mode_keys.ModeKeys.TRAIN:
+      clone._make_train_function()  # pylint: disable=protected-access
+      train_op = clone.train_function.updates_op
+    elif mode == mode_keys.ModeKeys.TEST:
+      clone._make_test_function()  # pylint: disable=protected-access
+    else:
+      clone._make_predict_function()  # pylint: disable=protected-access
+    g.get_collection_ref(ops.GraphKeys.UPDATE_OPS).extend(clone.state_updates)
+
+    with session.Session().as_default():
+      clone_var_list = _get_var_list(clone)
+      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 trackable equivalence.
+        _assert_same_non_optimizer_objects(model, model_graph, clone, g)
+
+        # TODO(b/113178242): Use value transfer for trackable 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  # pylint: disable=protected-access
+
+      # Add graph to the SavedModel builder.
+      builder.add_meta_graph(
+          model_utils.EXPORT_TAG_MAP[mode],
+          signature_def_map=_create_signature_def_map(clone, mode),
+          saver=saver_lib.Saver(
+              clone_var_list,
+              # Allow saving Models with no variables. This is somewhat odd, but
+              # it's not necessarily a bug.
+              allow_empty=True),
+          init_op=variables.local_variables_initializer(),
+          train_op=train_op)
+    return None
+
+
+def _create_signature_def_map(model, mode):
+  """Creates 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}  # pylint: disable=protected-access
+    inputs_dict.update(targets_dict)
+  outputs_dict = {name: x
+                  for name, x in zip(model.output_names, model.outputs)}
+  metrics = saving_utils.extract_model_metrics(model)
+
+  # Add metric variables to the `LOCAL_VARIABLES` collection. Metric variables
+  # are by default not added to any collections. We are doing this here, so
+  # that metric variables get initialized.
+  local_vars = set(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES))
+  vars_to_add = set()
+  if metrics is not None:
+    for key, value in metrics.items():
+      if isinstance(value, metrics_lib.Metric):
+        vars_to_add.update(value.variables)
+        # Convert Metric instances to (value_tensor, update_op) tuple.
+        metrics[key] = (value.result(), value.updates[0])
+  # Remove variables that are in the local variables collection already.
+  vars_to_add = vars_to_add.difference(local_vars)
+  for v in vars_to_add:
+    ops.add_to_collection(ops.GraphKeys.LOCAL_VARIABLES, v)
+
+  export_outputs = model_utils.export_outputs_for_mode(
+      mode,
+      predictions=outputs_dict,
+      loss=model.total_loss if model.optimizer else None,
+      metrics=metrics)
+  return model_utils.build_all_signature_defs(
+      inputs_dict,
+      export_outputs=export_outputs,
+      serving_only=(mode == mode_keys.ModeKeys.PREDICT))
+
+
+def _assert_same_non_optimizer_objects(model, model_graph, clone, clone_graph):  # pylint: disable=unused-argument
+  """Asserts model and clone contain the same trackable objects."""
+
+  # TODO(fchollet, kathywu): make sure this works in eager mode.
+  return True
+
+
+def load_from_saved_model(saved_model_path, custom_objects=None):
+  """Loads a keras Model from a SavedModel created by `export_saved_model()`.
+
+  This function reinstantiates model state by:
+  1) loading model topology from json (this will eventually come
+     from metagraph).
+  2) loading model weights from checkpoint.
+
+  Example:
+
+  ```python
+  import tensorflow as tf
+
+  # Create a tf.keras model.
+  model = tf.keras.Sequential()
+  model.add(tf.keras.layers.Dense(1, input_shape=[10]))
+  model.summary()
+
+  # Save the tf.keras model in the SavedModel format.
+  path = '/tmp/simple_keras_model'
+  tf.keras.experimental.export_saved_model(model, path)
+
+  # Load the saved keras model back.
+  new_model = tf.keras.experimental.load_from_saved_model(path)
+  new_model.summary()
+  ```
+
+  Args:
+    saved_model_path: a string specifying the path to an existing SavedModel.
+    custom_objects: Optional dictionary mapping names
+        (strings) to custom classes or functions to be
+        considered during deserialization.
+
+  Returns:
+    a keras.Model instance.
+  """
+  warnings.warn('`tf.keras.experimental.load_from_saved_model` is deprecated'
+                'and will be removed in a future version. '
+                'Please switch to `tf.keras.models.load_model`.')
+  # restore model topology from json string
+  model_json_filepath = os.path.join(
+      compat.as_bytes(saved_model_path),
+      compat.as_bytes(constants.ASSETS_DIRECTORY),
+      compat.as_bytes(SAVED_MODEL_FILENAME_JSON))
+  with gfile.Open(model_json_filepath, 'r') as f:
+    model_json = f.read()
+  model = model_config.model_from_json(
+      model_json, custom_objects=custom_objects)
+
+  # restore model weights
+  checkpoint_prefix = os.path.join(
+      compat.as_text(saved_model_path),
+      compat.as_text(constants.VARIABLES_DIRECTORY),
+      compat.as_text(constants.VARIABLES_FILENAME))
+  model.load_weights(checkpoint_prefix)
+  return model
+
+
+#### Directory / path helpers
+
+
+def _get_or_create_variables_dir(export_dir):
+  """Return variables sub-directory, or create one if it doesn't exist."""
+  variables_dir = _get_variables_dir(export_dir)
+  file_io.recursive_create_dir(variables_dir)
+  return variables_dir
+
+
+def _get_variables_dir(export_dir):
+  """Return variables sub-directory in the SavedModel."""
+  return os.path.join(
+      compat.as_text(export_dir),
+      compat.as_text(constants.VARIABLES_DIRECTORY))
+
+
+def _get_variables_path(export_dir):
+  """Return the variables path, used as the prefix for checkpoint files."""
+  return os.path.join(
+      compat.as_text(_get_variables_dir(export_dir)),
+      compat.as_text(constants.VARIABLES_FILENAME))
+
+
+def _get_or_create_assets_dir(export_dir):
+  """Return assets sub-directory, or create one if it doesn't exist."""
+  assets_destination_dir = _get_assets_dir(export_dir)
+
+  file_io.recursive_create_dir(assets_destination_dir)
+
+  return assets_destination_dir
+
+
+def _get_assets_dir(export_dir):
+  """Return path to asset directory in the SavedModel."""
+  return os.path.join(
+      compat.as_text(export_dir),
+      compat.as_text(constants.ASSETS_DIRECTORY))

tensorflow/python/keras/saving/saving_utils.py

@@ -0,0 +1,325 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils related to keras model saving."""
+
+import collections
+import copy
+import os
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras import losses
+from tensorflow.python.keras import optimizer_v1
+from tensorflow.python.keras import optimizers
+from tensorflow.python.keras.engine import base_layer_utils
+from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.keras.utils import version_utils
+from tensorflow.python.keras.utils.io_utils import ask_to_proceed_with_overwrite
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import nest
+
+
+def extract_model_metrics(model):
+  """Convert metrics from a Keras model `compile` API to dictionary.
+
+  This is used for converting Keras models to SavedModels.
+
+  Args:
+    model: A `tf.keras.Model` object.
+
+  Returns:
+    Dictionary mapping metric names to metric instances. May return `None` if
+    the model does not contain any metrics.
+  """
+  if getattr(model, '_compile_metrics', None):
+    # We are not using model.metrics here because we want to exclude the metrics
+    # added using `add_metric` API.
+    return {m.name: m for m in model._compile_metric_functions}  # pylint: disable=protected-access
+  return None
+
+
+def model_input_signature(model, keep_original_batch_size=False):
+  """Inspect model to get its input signature.
+
+  The model's input signature is a list with a single (possibly-nested) object.
+  This is due to the Keras-enforced restriction that tensor inputs must be
+  passed in as the first argument.
+
+  For example, a model with input {'feature1': <Tensor>, 'feature2': <Tensor>}
+  will have input signature: [{'feature1': TensorSpec, 'feature2': TensorSpec}]
+
+  Args:
+    model: Keras Model object.
+    keep_original_batch_size: A boolean indicating whether we want to keep using
+      the original batch size or set it to None. Default is `False`, which means
+      that the batch dim of the returned input signature will always be set to
+      `None`.
+
+  Returns:
+    A list containing either a single TensorSpec or an object with nested
+    TensorSpecs. This list does not contain the `training` argument.
+  """
+  input_specs = model._get_save_spec(dynamic_batch=not keep_original_batch_size)  # pylint: disable=protected-access
+  if input_specs is None:
+    return None
+  input_specs = _enforce_names_consistency(input_specs)
+  # Return a list with a single element as the model's input signature.
+  if isinstance(input_specs,
+                collections.abc.Sequence) and len(input_specs) == 1:
+    # Note that the isinstance check filters out single-element dictionaries,
+    # which should also be wrapped as a single-element list.
+    return input_specs
+  else:
+    return [input_specs]
+
+
+def raise_model_input_error(model):
+  raise ValueError(
+      'Model {} cannot be saved because the input shapes have not been '
+      'set. Usually, input shapes are automatically determined from calling'
+      ' `.fit()` or `.predict()`. To manually set the shapes, call '
+      '`model.build(input_shape)`.'.format(model))
+
+
+def trace_model_call(model, input_signature=None):
+  """Trace the model call to create a tf.function for exporting a Keras model.
+
+  Args:
+    model: A Keras model.
+    input_signature: optional, a list of tf.TensorSpec objects specifying the
+      inputs to the model.
+
+  Returns:
+    A tf.function wrapping the model's call function with input signatures set.
+
+  Raises:
+    ValueError: if input signature cannot be inferred from the model.
+  """
+  if input_signature is None:
+    if isinstance(model.call, def_function.Function):
+      input_signature = model.call.input_signature
+
+  if input_signature is None:
+    input_signature = model_input_signature(model)
+
+  if input_signature is None:
+    raise_model_input_error(model)
+
+  @def_function.function(input_signature=input_signature)
+  def _wrapped_model(*args):
+    """A concrete tf.function that wraps the model's call function."""
+    # When given a single input, Keras models will call the model on the tensor
+    # rather than a list consisting of the single tensor.
+    inputs = args[0] if len(input_signature) == 1 else list(args)
+
+    with base_layer_utils.call_context().enter(
+        model, inputs=inputs, build_graph=False, training=False, saving=True):
+      outputs = model(inputs, training=False)
+
+    # Outputs always has to be a flat dict.
+    output_names = model.output_names  # Functional Model.
+    if output_names is None:  # Subclassed Model.
+      from tensorflow.python.keras.engine import compile_utils  # pylint: disable=g-import-not-at-top
+      output_names = compile_utils.create_pseudo_output_names(outputs)
+    outputs = nest.flatten(outputs)
+    return {name: output for name, output in zip(output_names, outputs)}
+
+  return _wrapped_model
+
+
+def model_metadata(model, include_optimizer=True, require_config=True):
+  """Returns a dictionary containing the model metadata."""
+  from tensorflow.python.keras import __version__ as keras_version  # pylint: disable=g-import-not-at-top
+  from tensorflow.python.keras.optimizer_v2 import optimizer_v2  # pylint: disable=g-import-not-at-top
+
+  model_config = {'class_name': model.__class__.__name__}
+  try:
+    model_config['config'] = model.get_config()
+  except NotImplementedError as e:
+    if require_config:
+      raise e
+
+  metadata = dict(
+      keras_version=str(keras_version),
+      backend=K.backend(),
+      model_config=model_config)
+  if model.optimizer and include_optimizer:
+    if isinstance(model.optimizer, optimizer_v1.TFOptimizer):
+      logging.warning(
+          'TensorFlow optimizers do not '
+          'make it possible to access '
+          'optimizer attributes or optimizer state '
+          'after instantiation. '
+          'As a result, we cannot save the optimizer '
+          'as part of the model save file. '
+          'You will have to compile your model again after loading it. '
+          'Prefer using a Keras optimizer instead '
+          '(see keras.io/optimizers).')
+    elif model._compile_was_called:  # pylint: disable=protected-access
+      training_config = model._get_compile_args(user_metrics=False)  # pylint: disable=protected-access
+      training_config.pop('optimizer', None)  # Handled separately.
+      metadata['training_config'] = _serialize_nested_config(training_config)
+      if isinstance(model.optimizer, optimizer_v2.RestoredOptimizer):
+        raise NotImplementedError(
+            'As of now, Optimizers loaded from SavedModel cannot be saved. '
+            'If you\'re calling `model.save` or `tf.keras.models.save_model`,'
+            ' please set the `include_optimizer` option to `False`. For '
+            '`tf.saved_model.save`, delete the optimizer from the model.')
+      else:
+        optimizer_config = {
+            'class_name':
+                generic_utils.get_registered_name(model.optimizer.__class__),
+            'config':
+                model.optimizer.get_config()
+        }
+      metadata['training_config']['optimizer_config'] = optimizer_config
+  return metadata
+
+
+def should_overwrite(filepath, overwrite):
+  """Returns whether the filepath should be overwritten."""
+  # If file exists and should not be overwritten.
+  if not overwrite and os.path.isfile(filepath):
+    return ask_to_proceed_with_overwrite(filepath)
+  return True
+
+
+def compile_args_from_training_config(training_config, custom_objects=None):
+  """Return model.compile arguments from training config."""
+  if custom_objects is None:
+    custom_objects = {}
+
+  with generic_utils.CustomObjectScope(custom_objects):
+    optimizer_config = training_config['optimizer_config']
+    optimizer = optimizers.deserialize(optimizer_config)
+
+    # Recover losses.
+    loss = None
+    loss_config = training_config.get('loss', None)
+    if loss_config is not None:
+      loss = _deserialize_nested_config(losses.deserialize, loss_config)
+
+    # Recover metrics.
+    metrics = None
+    metrics_config = training_config.get('metrics', None)
+    if metrics_config is not None:
+      metrics = _deserialize_nested_config(_deserialize_metric, metrics_config)
+
+    # Recover weighted metrics.
+    weighted_metrics = None
+    weighted_metrics_config = training_config.get('weighted_metrics', None)
+    if weighted_metrics_config is not None:
+      weighted_metrics = _deserialize_nested_config(_deserialize_metric,
+                                                    weighted_metrics_config)
+
+    sample_weight_mode = training_config['sample_weight_mode'] if hasattr(
+        training_config, 'sample_weight_mode') else None
+    loss_weights = training_config['loss_weights']
+
+  return dict(
+      optimizer=optimizer,
+      loss=loss,
+      metrics=metrics,
+      weighted_metrics=weighted_metrics,
+      loss_weights=loss_weights,
+      sample_weight_mode=sample_weight_mode)
+
+
+def _deserialize_nested_config(deserialize_fn, config):
+  """Deserializes arbitrary Keras `config` using `deserialize_fn`."""
+
+  def _is_single_object(obj):
+    if isinstance(obj, dict) and 'class_name' in obj:
+      return True  # Serialized Keras object.
+    if isinstance(obj, str):
+      return True  # Serialized function or string.
+    return False
+
+  if config is None:
+    return None
+  if _is_single_object(config):
+    return deserialize_fn(config)
+  elif isinstance(config, dict):
+    return {
+        k: _deserialize_nested_config(deserialize_fn, v)
+        for k, v in config.items()
+    }
+  elif isinstance(config, (tuple, list)):
+    return [_deserialize_nested_config(deserialize_fn, obj) for obj in config]
+
+  raise ValueError('Saved configuration not understood.')
+
+
+def _serialize_nested_config(config):
+  """Serialized a nested structure of Keras objects."""
+
+  def _serialize_fn(obj):
+    if callable(obj):
+      return generic_utils.serialize_keras_object(obj)
+    return obj
+
+  return nest.map_structure(_serialize_fn, config)
+
+
+def _deserialize_metric(metric_config):
+  """Deserialize metrics, leaving special strings untouched."""
+  from tensorflow.python.keras import metrics as metrics_module  # pylint:disable=g-import-not-at-top
+  if metric_config in ['accuracy', 'acc', 'crossentropy', 'ce']:
+    # Do not deserialize accuracy and cross-entropy strings as we have special
+    # case handling for these in compile, based on model output shape.
+    return metric_config
+  return metrics_module.deserialize(metric_config)
+
+
+def _enforce_names_consistency(specs):
+  """Enforces that either all specs have names or none do."""
+
+  def _has_name(spec):
+    return hasattr(spec, 'name') and spec.name is not None
+
+  def _clear_name(spec):
+    spec = copy.deepcopy(spec)
+    if hasattr(spec, 'name'):
+      spec._name = None  # pylint:disable=protected-access
+    return spec
+
+  flat_specs = nest.flatten(specs)
+  name_inconsistency = (
+      any(_has_name(s) for s in flat_specs) and
+      not all(_has_name(s) for s in flat_specs))
+
+  if name_inconsistency:
+    specs = nest.map_structure(_clear_name, specs)
+  return specs
+
+
+def try_build_compiled_arguments(model):
+  if (not version_utils.is_v1_layer_or_model(model) and
+      model.outputs is not None):
+    try:
+      if not model.compiled_loss.built:
+        model.compiled_loss.build(model.outputs)
+      if not model.compiled_metrics.built:
+        model.compiled_metrics.build(model.outputs, model.outputs)
+    except:  # pylint: disable=bare-except
+      logging.warning(
+          'Compiled the loaded model, but the compiled metrics have yet to '
+          'be built. `model.compile_metrics` will be empty until you train '
+          'or evaluate the model.')
+
+
+def is_hdf5_filepath(filepath):
+  return (filepath.endswith('.h5') or filepath.endswith('.keras') or
+          filepath.endswith('.hdf5'))

tensorflow/python/keras/saving/utils_v1/BUILD

@@ -0,0 +1,55 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# Description:
+#   Keras saving and loading libraries.
+
+load("@xla//third_party/rules_python/python:defs.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/python/keras:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "utils_v1",
+    srcs = [
+        "__init__.py",
+        "export_output.py",
+        "export_utils.py",
+        "mode_keys.py",
+        "signature_def_utils.py",
+        "unexported_constants.py",
+    ],
+    srcs_version = "PY3",
+    strict_deps = False,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:signature_def_utils",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//tensorflow/python/saved_model:utils",
+        "//tensorflow/python/util:compat",
+    ],
+)

tensorflow/python/keras/saving/utils_v1/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# LINT.IfChange
+"""Utils for saving a Keras Model to the SavedModel format."""
+# pylint: disable=wildcard-import
+from tensorflow.python.keras.saving.utils_v1.export_output import *
+from tensorflow.python.keras.saving.utils_v1.export_utils import build_all_signature_defs
+from tensorflow.python.keras.saving.utils_v1.export_utils import export_outputs_for_mode
+from tensorflow.python.keras.saving.utils_v1.export_utils import EXPORT_TAG_MAP
+from tensorflow.python.keras.saving.utils_v1.export_utils import get_export_outputs
+from tensorflow.python.keras.saving.utils_v1.export_utils import get_temp_export_dir
+from tensorflow.python.keras.saving.utils_v1.export_utils import get_timestamped_export_dir
+from tensorflow.python.keras.saving.utils_v1.export_utils import SIGNATURE_KEY_MAP
+# pylint: enable=wildcard-import
+# LINT.ThenChange(//tensorflow/python/saved_model/model_utils/__init__.py)

tensorflow/python/keras/saving/utils_v1/export_output.py

@@ -0,0 +1,405 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# LINT.IfChange
+"""Classes for different types of export output."""
+
+import abc
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.keras.saving.utils_v1 import signature_def_utils as unexported_signature_utils
+from tensorflow.python.saved_model import signature_def_utils
+
+
+class ExportOutput(object):
+  """Represents an output of a model that can be served.
+
+  These typically correspond to model heads.
+  """
+
+  __metaclass__ = abc.ABCMeta
+
+  _SEPARATOR_CHAR = '/'
+
+  @abc.abstractmethod
+  def as_signature_def(self, receiver_tensors):
+    """Generate a SignatureDef proto for inclusion in a MetaGraphDef.
+
+    The SignatureDef will specify outputs as described in this ExportOutput,
+    and will use the provided receiver_tensors as inputs.
+
+    Args:
+      receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
+        input nodes that will be fed.
+    """
+    pass
+
+  def _check_output_key(self, key, error_label):
+    # For multi-head models, the key can be a tuple.
+    if isinstance(key, tuple):
+      key = self._SEPARATOR_CHAR.join(key)
+
+    if not isinstance(key, str):
+      raise ValueError(
+          '{} output key must be a string; got {}.'.format(error_label, key))
+    return key
+
+  def _wrap_and_check_outputs(
+      self, outputs, single_output_default_name, error_label=None):
+    """Wraps raw tensors as dicts and checks type.
+
+    Note that we create a new dict here so that we can overwrite the keys
+    if necessary.
+
+    Args:
+      outputs: A `Tensor` or a dict of string to `Tensor`.
+      single_output_default_name: A string key for use in the output dict
+        if the provided `outputs` is a raw tensor.
+      error_label: descriptive string for use in error messages. If none,
+        single_output_default_name will be used.
+
+    Returns:
+      A dict of tensors
+
+    Raises:
+      ValueError: if the outputs dict keys are not strings or tuples of strings
+        or the values are not Tensors.
+    """
+    if not isinstance(outputs, dict):
+      outputs = {single_output_default_name: outputs}
+
+    output_dict = {}
+    for key, value in outputs.items():
+      error_name = error_label or single_output_default_name
+      key = self._check_output_key(key, error_name)
+      if not isinstance(value, tensor.Tensor):
+        raise ValueError(
+            '{} output value must be a Tensor; got {}.'.format(
+                error_name, value))
+
+      output_dict[key] = value
+    return output_dict
+
+
+class ClassificationOutput(ExportOutput):
+  """Represents the output of a classification head.
+
+  Either classes or scores or both must be set.
+
+  The classes `Tensor` must provide string labels, not integer class IDs.
+
+  If only classes is set, it is interpreted as providing top-k results in
+  descending order.
+
+  If only scores is set, it is interpreted as providing a score for every class
+  in order of class ID.
+
+  If both classes and scores are set, they are interpreted as zipped, so each
+  score corresponds to the class at the same index.  Clients should not depend
+  on the order of the entries.
+  """
+
+  def __init__(self, scores=None, classes=None):
+    """Constructor for `ClassificationOutput`.
+
+    Args:
+      scores: A float `Tensor` giving scores (sometimes but not always
+          interpretable as probabilities) for each class.  May be `None`, but
+          only if `classes` is set.  Interpretation varies-- see class doc.
+      classes: A string `Tensor` giving predicted class labels.  May be `None`,
+          but only if `scores` is set.  Interpretation varies-- see class doc.
+
+    Raises:
+      ValueError: if neither classes nor scores is set, or one of them is not a
+          `Tensor` with the correct dtype.
+    """
+    if (scores is not None
+        and not (isinstance(scores, tensor.Tensor)
+                 and scores.dtype.is_floating)):
+      raise ValueError('Classification scores must be a float32 Tensor; '
+                       'got {}'.format(scores))
+    if (classes is not None
+        and not (isinstance(classes, tensor.Tensor)
+                 and dtypes.as_dtype(classes.dtype) == dtypes.string)):
+      raise ValueError('Classification classes must be a string Tensor; '
+                       'got {}'.format(classes))
+    if scores is None and classes is None:
+      raise ValueError('At least one of scores and classes must be set.')
+
+    self._scores = scores
+    self._classes = classes
+
+  @property
+  def scores(self):
+    return self._scores
+
+  @property
+  def classes(self):
+    return self._classes
+
+  def as_signature_def(self, receiver_tensors):
+    if len(receiver_tensors) != 1:
+      raise ValueError('Classification input must be a single string Tensor; '
+                       'got {}'.format(receiver_tensors))
+    (_, examples), = receiver_tensors.items()
+    if dtypes.as_dtype(examples.dtype) != dtypes.string:
+      raise ValueError('Classification input must be a single string Tensor; '
+                       'got {}'.format(receiver_tensors))
+    return signature_def_utils.classification_signature_def(
+        examples, self.classes, self.scores)
+
+
+class RegressionOutput(ExportOutput):
+  """Represents the output of a regression head."""
+
+  def __init__(self, value):
+    """Constructor for `RegressionOutput`.
+
+    Args:
+      value: a float `Tensor` giving the predicted values.  Required.
+
+    Raises:
+      ValueError: if the value is not a `Tensor` with dtype tf.float32.
+    """
+    if not (isinstance(value, tensor.Tensor) and value.dtype.is_floating):
+      raise ValueError('Regression output value must be a float32 Tensor; '
+                       'got {}'.format(value))
+    self._value = value
+
+  @property
+  def value(self):
+    return self._value
+
+  def as_signature_def(self, receiver_tensors):
+    if len(receiver_tensors) != 1:
+      raise ValueError('Regression input must be a single string Tensor; '
+                       'got {}'.format(receiver_tensors))
+    (_, examples), = receiver_tensors.items()
+    if dtypes.as_dtype(examples.dtype) != dtypes.string:
+      raise ValueError('Regression input must be a single string Tensor; '
+                       'got {}'.format(receiver_tensors))
+    return signature_def_utils.regression_signature_def(examples, self.value)
+
+
+class PredictOutput(ExportOutput):
+  """Represents the output of a generic prediction head.
+
+  A generic prediction need not be either a classification or a regression.
+
+  Named outputs must be provided as a dict from string to `Tensor`,
+  """
+  _SINGLE_OUTPUT_DEFAULT_NAME = 'output'
+
+  def __init__(self, outputs):
+    """Constructor for PredictOutput.
+
+    Args:
+      outputs: A `Tensor` or a dict of string to `Tensor` representing the
+        predictions.
+
+    Raises:
+      ValueError: if the outputs is not dict, or any of its keys are not
+          strings, or any of its values are not `Tensor`s.
+    """
+
+    self._outputs = self._wrap_and_check_outputs(
+        outputs, self._SINGLE_OUTPUT_DEFAULT_NAME, error_label='Prediction')
+
+  @property
+  def outputs(self):
+    return self._outputs
+
+  def as_signature_def(self, receiver_tensors):
+    return signature_def_utils.predict_signature_def(receiver_tensors,
+                                                     self.outputs)
+
+
+class _SupervisedOutput(ExportOutput):
+  """Represents the output of a supervised training or eval process."""
+  __metaclass__ = abc.ABCMeta
+
+  LOSS_NAME = 'loss'
+  PREDICTIONS_NAME = 'predictions'
+  METRICS_NAME = 'metrics'
+
+  METRIC_VALUE_SUFFIX = 'value'
+  METRIC_UPDATE_SUFFIX = 'update_op'
+
+  _loss = None
+  _predictions = None
+  _metrics = None
+
+  def __init__(self, loss=None, predictions=None, metrics=None):
+    """Constructor for SupervisedOutput (ie, Train or Eval output).
+
+    Args:
+      loss: dict of Tensors or single Tensor representing calculated loss.
+      predictions: dict of Tensors or single Tensor representing model
+        predictions.
+      metrics: Dict of metric results keyed by name.
+        The values of the dict can be one of the following:
+        (1) instance of `Metric` class.
+        (2) (metric_value, update_op) tuples, or a single tuple.
+        metric_value must be a Tensor, and update_op must be a Tensor or Op.
+
+    Raises:
+      ValueError: if any of the outputs' dict keys are not strings or tuples of
+        strings or the values are not Tensors (or Operations in the case of
+        update_op).
+    """
+
+    if loss is not None:
+      loss_dict = self._wrap_and_check_outputs(loss, self.LOSS_NAME)
+      self._loss = self._prefix_output_keys(loss_dict, self.LOSS_NAME)
+    if predictions is not None:
+      pred_dict = self._wrap_and_check_outputs(
+          predictions, self.PREDICTIONS_NAME)
+      self._predictions = self._prefix_output_keys(
+          pred_dict, self.PREDICTIONS_NAME)
+    if metrics is not None:
+      self._metrics = self._wrap_and_check_metrics(metrics)
+
+  def _prefix_output_keys(self, output_dict, output_name):
+    """Prepend output_name to the output_dict keys if it doesn't exist.
+
+    This produces predictable prefixes for the pre-determined outputs
+    of SupervisedOutput.
+
+    Args:
+      output_dict: dict of string to Tensor, assumed valid.
+      output_name: prefix string to prepend to existing keys.
+
+    Returns:
+      dict with updated keys and existing values.
+    """
+
+    new_outputs = {}
+    for key, val in output_dict.items():
+      key = self._prefix_key(key, output_name)
+      new_outputs[key] = val
+    return new_outputs
+
+  def _prefix_key(self, key, output_name):
+    if key.find(output_name) != 0:
+      key = output_name + self._SEPARATOR_CHAR + key
+    return key
+
+  def _wrap_and_check_metrics(self, metrics):
+    """Handle the saving of metrics.
+
+    Metrics is either a tuple of (value, update_op), or a dict of such tuples.
+    Here, we separate out the tuples and create a dict with names to tensors.
+
+    Args:
+      metrics: Dict of metric results keyed by name.
+        The values of the dict can be one of the following:
+        (1) instance of `Metric` class.
+        (2) (metric_value, update_op) tuples, or a single tuple.
+        metric_value must be a Tensor, and update_op must be a Tensor or Op.
+
+    Returns:
+      dict of output_names to tensors
+
+    Raises:
+      ValueError: if the dict key is not a string, or the metric values or ops
+        are not tensors.
+    """
+    if not isinstance(metrics, dict):
+      metrics = {self.METRICS_NAME: metrics}
+
+    outputs = {}
+    for key, value in metrics.items():
+      if isinstance(value, tuple):
+        metric_val, metric_op = value
+      else:  # value is a keras.Metrics object
+        metric_val = value.result()
+        assert len(value.updates) == 1  # We expect only one update op.
+        metric_op = value.updates[0]
+      key = self._check_output_key(key, self.METRICS_NAME)
+      key = self._prefix_key(key, self.METRICS_NAME)
+
+      val_name = key + self._SEPARATOR_CHAR + self.METRIC_VALUE_SUFFIX
+      op_name = key + self._SEPARATOR_CHAR + self.METRIC_UPDATE_SUFFIX
+      if not isinstance(metric_val, tensor.Tensor):
+        raise ValueError(
+            '{} output value must be a Tensor; got {}.'.format(
+                key, metric_val))
+      if not (tensor_util.is_tensor(metric_op) or
+              isinstance(metric_op, ops.Operation)):
+        raise ValueError(
+            '{} update_op must be a Tensor or Operation; got {}.'.format(
+                key, metric_op))
+
+      # We must wrap any ops (or variables) in a Tensor before export, as the
+      # SignatureDef proto expects tensors only. See b/109740581
+      metric_op_tensor = metric_op
+      if not isinstance(metric_op, tensor.Tensor):
+        with ops.control_dependencies([metric_op]):
+          metric_op_tensor = constant_op.constant([], name='metric_op_wrapper')
+
+      outputs[val_name] = metric_val
+      outputs[op_name] = metric_op_tensor
+
+    return outputs
+
+  @property
+  def loss(self):
+    return self._loss
+
+  @property
+  def predictions(self):
+    return self._predictions
+
+  @property
+  def metrics(self):
+    return self._metrics
+
+  @abc.abstractmethod
+  def _get_signature_def_fn(self):
+    """Returns a function that produces a SignatureDef given desired outputs."""
+    pass
+
+  def as_signature_def(self, receiver_tensors):
+    signature_def_fn = self._get_signature_def_fn()
+    return signature_def_fn(
+        receiver_tensors, self.loss, self.predictions, self.metrics)
+
+
+class TrainOutput(_SupervisedOutput):
+  """Represents the output of a supervised training process.
+
+  This class generates the appropriate signature def for exporting
+  training output by type-checking and wrapping loss, predictions, and metrics
+  values.
+  """
+
+  def _get_signature_def_fn(self):
+    return unexported_signature_utils.supervised_train_signature_def
+
+
+class EvalOutput(_SupervisedOutput):
+  """Represents the output of a supervised eval process.
+
+  This class generates the appropriate signature def for exporting
+  eval output by type-checking and wrapping loss, predictions, and metrics
+  values.
+  """
+
+  def _get_signature_def_fn(self):
+    return unexported_signature_utils.supervised_eval_signature_def
+# LINT.ThenChange(//tensorflow/python/saved_model/model_utils/export_output.py)

tensorflow/python/keras/saving/utils_v1/export_utils.py

@@ -0,0 +1,356 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# LINT.IfChange
+"""Utilities for creating SavedModels."""
+
+import collections
+import os
+import time
+
+from tensorflow.python.keras.saving.utils_v1 import export_output as export_output_lib
+from tensorflow.python.keras.saving.utils_v1 import mode_keys
+from tensorflow.python.keras.saving.utils_v1 import unexported_constants
+from tensorflow.python.keras.saving.utils_v1.mode_keys import KerasModeKeys as ModeKeys
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.util import compat
+
+
+# Mapping of the modes to appropriate MetaGraph tags in the SavedModel.
+EXPORT_TAG_MAP = mode_keys.ModeKeyMap(**{
+    ModeKeys.PREDICT: [tag_constants.SERVING],
+    ModeKeys.TRAIN: [tag_constants.TRAINING],
+    ModeKeys.TEST: [unexported_constants.EVAL]})
+
+# For every exported mode, a SignatureDef map should be created using the
+# functions `export_outputs_for_mode` and `build_all_signature_defs`. By
+# default, this map will contain a single Signature that defines the input
+# tensors and output predictions, losses, and/or metrics (depending on the mode)
+# The default keys used in the SignatureDef map are defined below.
+SIGNATURE_KEY_MAP = mode_keys.ModeKeyMap(**{
+    ModeKeys.PREDICT: signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
+    ModeKeys.TRAIN: unexported_constants.DEFAULT_TRAIN_SIGNATURE_DEF_KEY,
+    ModeKeys.TEST: unexported_constants.DEFAULT_EVAL_SIGNATURE_DEF_KEY})
+
+# Default names used in the SignatureDef input map, which maps strings to
+# TensorInfo protos.
+SINGLE_FEATURE_DEFAULT_NAME = 'feature'
+SINGLE_RECEIVER_DEFAULT_NAME = 'input'
+SINGLE_LABEL_DEFAULT_NAME = 'label'
+
+### Below utilities are specific to SavedModel exports.
+
+
+def build_all_signature_defs(receiver_tensors,
+                             export_outputs,
+                             receiver_tensors_alternatives=None,
+                             serving_only=True):
+  """Build `SignatureDef`s for all export outputs.
+
+  Args:
+    receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
+      input nodes where this receiver expects to be fed by default.  Typically,
+      this is a single placeholder expecting serialized `tf.Example` protos.
+    export_outputs: a dict of ExportOutput instances, each of which has
+      an as_signature_def instance method that will be called to retrieve
+      the signature_def for all export output tensors.
+    receiver_tensors_alternatives: a dict of string to additional
+      groups of receiver tensors, each of which may be a `Tensor` or a dict of
+      string to `Tensor`.  These named receiver tensor alternatives generate
+      additional serving signatures, which may be used to feed inputs at
+      different points within the input receiver subgraph.  A typical usage is
+      to allow feeding raw feature `Tensor`s *downstream* of the
+      tf.io.parse_example() op.  Defaults to None.
+    serving_only: boolean; if true, resulting signature defs will only include
+      valid serving signatures. If false, all requested signatures will be
+      returned.
+
+  Returns:
+    signature_def representing all passed args.
+
+  Raises:
+    ValueError: if export_outputs is not a dict
+  """
+  if not isinstance(receiver_tensors, dict):
+    receiver_tensors = {SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors}
+  if export_outputs is None or not isinstance(export_outputs, dict):
+    raise ValueError('export_outputs must be a dict and not'
+                     '{}'.format(type(export_outputs)))
+
+  signature_def_map = {}
+  excluded_signatures = {}
+  for output_key, export_output in export_outputs.items():
+    signature_name = '{}'.format(output_key or 'None')
+    try:
+      signature = export_output.as_signature_def(receiver_tensors)
+      signature_def_map[signature_name] = signature
+    except ValueError as e:
+      excluded_signatures[signature_name] = str(e)
+
+  if receiver_tensors_alternatives:
+    for receiver_name, receiver_tensors_alt in (
+        receiver_tensors_alternatives.items()):
+      if not isinstance(receiver_tensors_alt, dict):
+        receiver_tensors_alt = {
+            SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors_alt
+        }
+      for output_key, export_output in export_outputs.items():
+        signature_name = '{}:{}'.format(receiver_name or 'None', output_key or
+                                        'None')
+        try:
+          signature = export_output.as_signature_def(receiver_tensors_alt)
+          signature_def_map[signature_name] = signature
+        except ValueError as e:
+          excluded_signatures[signature_name] = str(e)
+
+  _log_signature_report(signature_def_map, excluded_signatures)
+
+  # The above calls to export_output_lib.as_signature_def should return only
+  # valid signatures; if there is a validity problem, they raise a ValueError,
+  # in which case we exclude that signature from signature_def_map above.
+  # The is_valid_signature check ensures that the signatures produced are
+  # valid for serving, and acts as an additional sanity check for export
+  # signatures produced for serving. We skip this check for training and eval
+  # signatures, which are not intended for serving.
+  if serving_only:
+    signature_def_map = {
+        k: v
+        for k, v in signature_def_map.items()
+        if signature_def_utils.is_valid_signature(v)
+    }
+  return signature_def_map
+
+
+_FRIENDLY_METHOD_NAMES = {
+    signature_constants.CLASSIFY_METHOD_NAME: 'Classify',
+    signature_constants.REGRESS_METHOD_NAME: 'Regress',
+    signature_constants.PREDICT_METHOD_NAME: 'Predict',
+    unexported_constants.SUPERVISED_TRAIN_METHOD_NAME: 'Train',
+    unexported_constants.SUPERVISED_EVAL_METHOD_NAME: 'Eval',
+}
+
+
+def _log_signature_report(signature_def_map, excluded_signatures):
+  """Log a report of which signatures were produced."""
+  sig_names_by_method_name = collections.defaultdict(list)
+
+  # We'll collect whatever method_names are present, but also we want to make
+  # sure to output a line for each of the three standard methods even if they
+  # have no signatures.
+  for method_name in _FRIENDLY_METHOD_NAMES:
+    sig_names_by_method_name[method_name] = []
+
+  for signature_name, sig in signature_def_map.items():
+    sig_names_by_method_name[sig.method_name].append(signature_name)
+
+  # TODO(b/67733540): consider printing the full signatures, not just names
+  for method_name, sig_names in sig_names_by_method_name.items():
+    if method_name in _FRIENDLY_METHOD_NAMES:
+      method_name = _FRIENDLY_METHOD_NAMES[method_name]
+    logging.info('Signatures INCLUDED in export for {}: {}'.format(
+        method_name, sig_names if sig_names else 'None'))
+
+  if excluded_signatures:
+    logging.info('Signatures EXCLUDED from export because they cannot be '
+                 'be served via TensorFlow Serving APIs:')
+    for signature_name, message in excluded_signatures.items():
+      logging.info('\'{}\' : {}'.format(signature_name, message))
+
+  if not signature_def_map:
+    logging.warning('Export includes no signatures!')
+  elif (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY not in
+        signature_def_map):
+    logging.warning('Export includes no default signature!')
+
+
+# When we create a timestamped directory, there is a small chance that the
+# directory already exists because another process is also creating these
+# directories. In this case we just wait one second to get a new timestamp and
+# try again. If this fails several times in a row, then something is seriously
+# wrong.
+MAX_DIRECTORY_CREATION_ATTEMPTS = 10
+
+
+def get_timestamped_export_dir(export_dir_base):
+  """Builds a path to a new subdirectory within the base directory.
+
+  Each export is written into a new subdirectory named using the
+  current time.  This guarantees monotonically increasing version
+  numbers even across multiple runs of the pipeline.
+  The timestamp used is the number of seconds since epoch UTC.
+
+  Args:
+    export_dir_base: A string containing a directory to write the exported
+        graph and checkpoints.
+  Returns:
+    The full path of the new subdirectory (which is not actually created yet).
+
+  Raises:
+    RuntimeError: if repeated attempts fail to obtain a unique timestamped
+      directory name.
+  """
+  attempts = 0
+  while attempts < MAX_DIRECTORY_CREATION_ATTEMPTS:
+    timestamp = int(time.time())
+
+    result_dir = os.path.join(
+        compat.as_bytes(export_dir_base), compat.as_bytes(str(timestamp)))
+    if not gfile.Exists(result_dir):
+      # Collisions are still possible (though extremely unlikely): this
+      # directory is not actually created yet, but it will be almost
+      # instantly on return from this function.
+      return result_dir
+    time.sleep(1)
+    attempts += 1
+    logging.warning(
+        'Directory {} already exists; retrying (attempt {}/{})'.format(
+            compat.as_str(result_dir), attempts,
+            MAX_DIRECTORY_CREATION_ATTEMPTS))
+  raise RuntimeError('Failed to obtain a unique export directory name after '
+                     '{} attempts.'.format(MAX_DIRECTORY_CREATION_ATTEMPTS))
+
+
+def get_temp_export_dir(timestamped_export_dir):
+  """Builds a directory name based on the argument but starting with 'temp-'.
+
+  This relies on the fact that TensorFlow Serving ignores subdirectories of
+  the base directory that can't be parsed as integers.
+
+  Args:
+    timestamped_export_dir: the name of the eventual export directory, e.g.
+      /foo/bar/<timestamp>
+
+  Returns:
+    A sister directory prefixed with 'temp-', e.g. /foo/bar/temp-<timestamp>.
+  """
+  (dirname, basename) = os.path.split(timestamped_export_dir)
+  if isinstance(basename, bytes):
+    str_name = basename.decode('utf-8')
+  else:
+    str_name = str(basename)
+  temp_export_dir = os.path.join(
+      compat.as_bytes(dirname),
+      compat.as_bytes('temp-{}'.format(str_name)))
+  return temp_export_dir
+
+
+def export_outputs_for_mode(
+    mode, serving_export_outputs=None, predictions=None, loss=None,
+    metrics=None):
+  """Util function for constructing a `ExportOutput` dict given a mode.
+
+  The returned dict can be directly passed to `build_all_signature_defs` helper
+  function as the `export_outputs` argument, used for generating a SignatureDef
+  map.
+
+  Args:
+    mode: A `ModeKeys` specifying the mode.
+    serving_export_outputs: Describes the output signatures to be exported to
+      `SavedModel` and used during serving. Should be a dict or None.
+    predictions: A dict of Tensors or single Tensor representing model
+        predictions. This argument is only used if serving_export_outputs is not
+        set.
+    loss: A dict of Tensors or single Tensor representing calculated loss.
+    metrics: A dict of (metric_value, update_op) tuples, or a single tuple.
+      metric_value must be a Tensor, and update_op must be a Tensor or Op
+
+  Returns:
+    Dictionary mapping the key to an `ExportOutput` object.
+    The key is the expected SignatureDef key for the mode.
+
+  Raises:
+    ValueError: if an appropriate ExportOutput cannot be found for the mode.
+  """
+  if mode not in SIGNATURE_KEY_MAP:
+    raise ValueError(
+        'Export output type not found for mode: {}. Expected one of: {}.\n'
+        .format(mode, SIGNATURE_KEY_MAP.keys()))
+  signature_key = SIGNATURE_KEY_MAP[mode]
+  if mode_keys.is_predict(mode):
+    return get_export_outputs(serving_export_outputs, predictions)
+  elif mode_keys.is_train(mode):
+    return {signature_key: export_output_lib.TrainOutput(
+        loss=loss, predictions=predictions, metrics=metrics)}
+  else:
+    return {signature_key: export_output_lib.EvalOutput(
+        loss=loss, predictions=predictions, metrics=metrics)}
+
+
+def get_export_outputs(export_outputs, predictions):
+  """Validate export_outputs or create default export_outputs.
+
+  Args:
+    export_outputs: Describes the output signatures to be exported to
+      `SavedModel` and used during serving. Should be a dict or None.
+    predictions:  Predictions `Tensor` or dict of `Tensor`.
+
+  Returns:
+    Valid export_outputs dict
+
+  Raises:
+    TypeError: if export_outputs is not a dict or its values are not
+      ExportOutput instances.
+  """
+  if export_outputs is None:
+    default_output = export_output_lib.PredictOutput(predictions)
+    export_outputs = {
+        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: default_output}
+
+  if not isinstance(export_outputs, dict):
+    raise TypeError('export_outputs must be dict, given: {}'.format(
+        export_outputs))
+  for v in export_outputs.values():
+    if not isinstance(v, export_output_lib.ExportOutput):
+      raise TypeError(
+          'Values in export_outputs must be ExportOutput objects. '
+          'Given: {}'.format(export_outputs))
+
+  _maybe_add_default_serving_output(export_outputs)
+
+  return export_outputs
+
+
+def _maybe_add_default_serving_output(export_outputs):
+  """Add a default serving output to the export_outputs if not present.
+
+  Args:
+    export_outputs: Describes the output signatures to be exported to
+      `SavedModel` and used during serving. Should be a dict.
+
+  Returns:
+    export_outputs dict with default serving signature added if necessary
+
+  Raises:
+    ValueError: if multiple export_outputs were provided without a default
+      serving key.
+  """
+  if len(export_outputs) == 1:
+    (key, value), = export_outputs.items()
+    if key != signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+      export_outputs[
+          signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = value
+  if len(export_outputs) > 1:
+    if (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+        not in export_outputs):
+      raise ValueError(
+          'Multiple export_outputs were provided, but none of them is '
+          'specified as the default.  Do this by naming one of them with '
+          'signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY.')
+
+  return export_outputs
+# LINT.ThenChange(//tensorflow/python/saved_model/model_utils/export_utils.py)

tensorflow/python/keras/saving/utils_v1/mode_keys.py

@@ -0,0 +1,87 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# LINT.IfChange
+"""Utils for managing different mode strings used by Keras models.
+"""
+
+import collections
+
+
+class KerasModeKeys:
+  """Standard names for model modes.
+
+  The following standard keys are defined:
+
+  * `TRAIN`: training/fitting mode.
+  * `TEST`: testing/evaluation mode.
+  * `PREDICT`: prediction/inference mode.
+  """
+
+  TRAIN = 'train'
+  TEST = 'test'
+  PREDICT = 'predict'
+
+
+def is_predict(mode):
+  return mode == KerasModeKeys.PREDICT
+
+
+def is_eval(mode):
+  return mode == KerasModeKeys.TEST
+
+
+def is_train(mode):
+  return mode == KerasModeKeys.TRAIN
+
+
+class ModeKeyMap(collections.abc.Mapping):
+  """Map using ModeKeys as keys.
+
+  This class creates an immutable mapping from modes to values. For example,
+  SavedModel export of Keras models use this to map modes to their
+  corresponding MetaGraph tags/SignatureDef keys.
+  """
+
+  def __init__(self, **kwargs):
+    self._internal_dict = {}
+    self._keys = []
+    for key in kwargs:
+      self._keys.append(key)
+      dict_key = self._get_internal_key(key)
+      if dict_key in self._internal_dict:
+        raise ValueError(
+            'Error creating ModeKeyMap. Multiple keys/values found for {} mode.'
+            .format(dict_key))
+      self._internal_dict[dict_key] = kwargs[key]
+
+  def _get_internal_key(self, key):
+    """Return keys used for the internal dictionary."""
+    if is_train(key):
+      return KerasModeKeys.TRAIN
+    if is_eval(key):
+      return KerasModeKeys.TEST
+    if is_predict(key):
+      return KerasModeKeys.PREDICT
+    raise ValueError('Invalid mode key: {}.'.format(key))
+
+  def __getitem__(self, key):
+    return self._internal_dict[self._get_internal_key(key)]
+
+  def __iter__(self):
+    return iter(self._keys)
+
+  def __len__(self):
+    return len(self._keys)
+# LINT.ThenChange(//tensorflow/python/saved_model/model_utils/mode_keys.py)

tensorflow/python/keras/saving/utils_v1/signature_def_utils.py

@@ -0,0 +1,77 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SignatureDef utility functions implementation."""
+
+from tensorflow.python.keras.saving.utils_v1 import unexported_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import utils_impl as utils
+
+
+# LINT.IfChange
+def supervised_train_signature_def(
+    inputs, loss, predictions=None, metrics=None):
+  return _supervised_signature_def(
+      unexported_constants.SUPERVISED_TRAIN_METHOD_NAME, inputs, loss=loss,
+      predictions=predictions, metrics=metrics)
+
+
+def supervised_eval_signature_def(
+    inputs, loss, predictions=None, metrics=None):
+  return _supervised_signature_def(
+      unexported_constants.SUPERVISED_EVAL_METHOD_NAME, inputs, loss=loss,
+      predictions=predictions, metrics=metrics)
+
+
+def _supervised_signature_def(
+    method_name, inputs, loss=None, predictions=None,
+    metrics=None):
+  """Creates a signature for training and eval data.
+
+  This function produces signatures that describe the inputs and outputs
+  of a supervised process, such as training or evaluation, that
+  results in loss, metrics, and the like. Note that this function only requires
+  inputs to be not None.
+
+  Args:
+    method_name: Method name of the SignatureDef as a string.
+    inputs: dict of string to `Tensor`.
+    loss: dict of string to `Tensor` representing computed loss.
+    predictions: dict of string to `Tensor` representing the output predictions.
+    metrics: dict of string to `Tensor` representing metric ops.
+
+  Returns:
+    A train- or eval-flavored signature_def.
+
+  Raises:
+    ValueError: If inputs or outputs is `None`.
+  """
+  if inputs is None or not inputs:
+    raise ValueError('{} inputs cannot be None or empty.'.format(method_name))
+
+  signature_inputs = {key: utils.build_tensor_info(tensor)
+                      for key, tensor in inputs.items()}
+
+  signature_outputs = {}
+  for output_set in (loss, predictions, metrics):
+    if output_set is not None:
+      sig_out = {key: utils.build_tensor_info(tensor)
+                 for key, tensor in output_set.items()}
+      signature_outputs.update(sig_out)
+
+  signature_def = signature_def_utils.build_signature_def(
+      signature_inputs, signature_outputs, method_name)
+
+  return signature_def
+# LINT.ThenChange(//tensorflow/python/keras/saving/utils_v1/signature_def_utils.py)

tensorflow/python/keras/saving/utils_v1/unexported_constants.py

@@ -0,0 +1,32 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Signature constants for SavedModel save and restore operations.
+
+These are the private constants that have not been exported.
+"""
+
+# LINT.IfChange
+DEFAULT_TRAIN_SIGNATURE_DEF_KEY = "train"
+
+DEFAULT_EVAL_SIGNATURE_DEF_KEY = "eval"
+
+SUPERVISED_TRAIN_METHOD_NAME = "tensorflow/supervised/training"
+
+SUPERVISED_EVAL_METHOD_NAME = "tensorflow/supervised/eval"
+# LINT.ThenChange(//tensorflow/python/saved_model/signature_constants.py)
+
+# LINT.IfChange
+EVAL = "eval"
+# LINT.ThenChange(//tensorflow/python/saved_model/tag_constants.py)