No-op refactor of meta graph editing into its own internal library.

PiperOrigin-RevId: 212772668

tensorflow_hub/BUILD

@@ -197,6 +197,7 @@ py_library(
     srcs_version = "PY2AND3",
     deps = [
         ":all_protos_py_pb2",
+        ":meta_graph_lib",
         ":module",
         ":saved_model_lib",
         ":tf_utils",
@@ -283,6 +284,16 @@ py_test(
     ],
 )
 
+py_library(
+    name = "meta_graph_lib",
+    srcs = ["meta_graph_lib.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//tensorflow_hub:expect_tensorflow_installed",
+        #  "tensorflow/core:protos_all_py_pb2",
+    ],
+)
+
 py_library(
     name = "tensor_info",
     srcs = ["tensor_info.py"],

tensorflow_hub/meta_graph_lib.py

@@ -0,0 +1,61 @@
+# Copyright 2018 The TensorFlow Hub 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.
+# ==============================================================================
+"""MetaGraph lib provides utilities to manipulate MetaGraphDefs.
+
+This is an internal Hub utility and not part of the public API.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import re
+
+import tensorflow as tf
+
+
+def prepend_name_scope(name, import_scope):
+  """Prepends name scope to a name."""
+  # Based on tensorflow/python/framework/ops.py implementation.
+  if import_scope:
+    try:
+      str_to_replace = r"([\^]|loc:@|^)(.*)"
+      return re.sub(str_to_replace, r"\1" + import_scope + r"/\2",
+                    tf.compat.as_str_any(name))
+    except TypeError as e:
+      # If the name is not of a type we can process, simply return it.
+      tf.logging.warning(e)
+      return name
+  else:
+    return name
+
+
+def prefix_shared_name_attributes(meta_graph, absolute_import_scope):
+  """In-place prefixes shared_name attributes of nodes."""
+  shared_name_attr = "shared_name"
+  for node in meta_graph.graph_def.node:
+    shared_name_value = node.attr.get(shared_name_attr, None)
+    if shared_name_value and shared_name_value.HasField("s"):
+      if shared_name_value.s:
+        node.attr[shared_name_attr].s = tf.compat.as_bytes(
+            prepend_name_scope(
+                shared_name_value.s, import_scope=absolute_import_scope))
+
+
+def filter_collections(meta_graph, collections):
+  collections = frozenset(collections)
+  for name in list(meta_graph.collection_def.keys()):
+    if name not in collections:
+      del meta_graph.collection_def[name]

tensorflow_hub/native_module.py

@@ -24,6 +24,7 @@
 
 import tensorflow as tf
 
+from tensorflow_hub import meta_graph_lib
 from tensorflow_hub import module_def_pb2
 from tensorflow_hub import module_impl
 from tensorflow_hub import module_spec
@@ -433,10 +434,17 @@ def _create_state_graph(self, name):
     relative_scope_name = absolute_scope_name.split("/")[-1]
     assert relative_scope_name == name  # verify name scope was indeed unused.
 
+    meta_graph = meta_graph_pb2.MetaGraphDef()
+    meta_graph.CopyFrom(self._meta_graph)
+
+    meta_graph_lib.filter_collections(meta_graph, import_collections)
+    meta_graph_lib.prefix_shared_name_attributes(meta_graph,
+                                                 absolute_scope_name)
+
     tf.train.import_meta_graph(
-        adapted_meta_graph_for_import(self._meta_graph, absolute_scope_name),
-        import_scope=relative_scope_name,
-        restore_collections_predicate=(lambda key: key in import_collections))
+        meta_graph,
+        input_map={},
+        import_scope=relative_scope_name)
 
     # Build a list from the variable name in the module definition to the actual
     # instantiated variables.
@@ -449,7 +457,8 @@ def _create_state_graph(self, name):
     # apply-graphs.
     def _get_tensor(tensor_name):
       return self._graph.get_tensor_by_name(
-          prepend_name_scope(tensor_name, import_scope=absolute_scope_name))
+          meta_graph_lib.prepend_name_scope(tensor_name,
+                                            import_scope=absolute_scope_name))
 
     state_op_names = list_registered_stateful_ops_without_inputs()
     state_map = get_state_map(self._meta_graph, state_op_names, set(),
@@ -496,11 +505,17 @@ def create_apply_graph(self, signature, inputs, name):
     if self._trainable:
       import_collections.extend([tf.GraphKeys.UPDATE_OPS])
 
+    meta_graph = meta_graph_pb2.MetaGraphDef()
+    meta_graph.CopyFrom(self._meta_graph)
+
+    meta_graph_lib.filter_collections(meta_graph, import_collections)
+    meta_graph_lib.prefix_shared_name_attributes(meta_graph,
+                                                 absolute_scope_name)
+
     tf.train.import_meta_graph(
-        adapted_meta_graph_for_import(self._meta_graph, absolute_scope_name),
+        meta_graph,
         input_map=feed_map,
-        import_scope=relative_scope_name,
-        restore_collections_predicate=(lambda key: key in import_collections))
+        import_scope=relative_scope_name)
     fix_colocation_after_import(input_map=feed_map,
                                 absolute_import_scope=absolute_scope_name)
 
@@ -511,7 +526,8 @@ def get_tensor(name):
         return feed_map[name]
       except KeyError:
         return self._graph.get_tensor_by_name(
-            prepend_name_scope(name, import_scope=absolute_scope_name))
+            meta_graph_lib.prepend_name_scope(name,
+                                              import_scope=absolute_scope_name))
 
     return tensor_info.build_output_map(signature_def.outputs, get_tensor)
 
@@ -532,39 +548,6 @@ def variable_map(self):
     return self._variable_map
 
 
-def adapted_meta_graph_for_import(meta_graph, absolute_import_scope):
-  """Prefixes the shared_name attributes of nodes with the import scope."""
-  # The MetaGraphDef is copied (and not modified in-place) so that future
-  # imports won't see these modifications.
-  copy = meta_graph_pb2.MetaGraphDef()
-  copy.CopyFrom(meta_graph)
-  shared_name_attr = "shared_name"
-  for node in copy.graph_def.node:
-    shared_name_value = node.attr.get(shared_name_attr, None)
-    if shared_name_value and shared_name_value.HasField("s"):
-      if shared_name_value.s:
-        node.attr[shared_name_attr].s = tf.compat.as_bytes(
-            prepend_name_scope(
-                shared_name_value.s, import_scope=absolute_import_scope))
-  return copy
-
-
-def prepend_name_scope(name, import_scope):
-  """Prepends name scope to a name."""
-  # Based on tensorflow/python/framework/ops.py implementation.
-  if import_scope:
-    try:
-      str_to_replace = r"([\^]|loc:@|^)(.*)"
-      return re.sub(str_to_replace, r"\1" + import_scope + r"/\2",
-                    tf.compat.as_str_any(name))
-    except TypeError as e:
-      # If the name is not of a type we can process, simply return it.
-      tf.logging.warning(e)
-      return name
-  else:
-    return name
-
-
 def list_registered_stateful_ops_without_inputs():
   """Returns set of registered stateful ops that do not expect inputs.