Add SavedModel.LoadOptions to hub.KerasLayer API to pass to load_v2.

PiperOrigin-RevId: 322425049
Change-Id: I4fd626c4a7e470bcb647a3d76d1ece16fcb63b28

tensorflow/python/keras/saving/save.py

@@ -27,6 +27,7 @@ from tensorflow.python.keras.saving.saved_model import load as saved_model_load
 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
+from tensorflow.python.saved_model import load_context
 from tensorflow.python.saved_model import loader_impl
 from tensorflow.python.util.tf_export import keras_export
 
@@ -177,14 +178,16 @@ def load_model(filepath, custom_objects=None, compile=True, options=None):  # py
       IOError: In case of an invalid savefile.
   """
   with generic_utils.CustomObjectScope(custom_objects or {}):
-    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)
+    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, six.string_types):
-      loader_impl.parse_saved_model(filepath)
-      return saved_model_load.load(filepath, compile, options)
+      filepath = path_to_string(filepath)
+      if isinstance(filepath, six.string_types):
+        loader_impl.parse_saved_model(filepath)
+        return saved_model_load.load(filepath, compile, options)
 
   raise IOError(
       'Unable to load model. Filepath is not an hdf5 file (or h5py is not '

tensorflow/python/saved_model/BUILD

@@ -380,6 +380,15 @@ tf_py_test(
 )
 
 py_strict_library(
+    name = "load_context",
+    srcs = [
+        "load_context.py",
+    ],
+    srcs_version = "PY2AND3",
+    deps = [],
+)
+
+py_library(
     name = "load",
     srcs = [
         "load.py",
@@ -387,6 +396,7 @@ py_strict_library(
     srcs_version = "PY2AND3",
     deps = [
         ":function_deserialization",
+        ":load_context",
         ":load_options",
         ":load_v1_in_v2",
         ":loader",

tensorflow/python/saved_model/load_context.py

@@ -0,0 +1,56 @@
+# 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."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import threading
+
+
+class LoadContext(threading.local):
+  """A context for loading a model."""
+
+  def __init__(self):
+    super(LoadContext, self).__init__()
+    self._load_options = None
+
+  def set_load_options(self, load_options):
+    self._load_options = load_options
+
+  def clear_load_options(self):
+    self._load_options = None
+
+  def load_options(self):
+    return self._load_options
+
+
+_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 whether under a load context."""
+  return _load_context.load_options()