Fixed two bugs when importing MetaGraphDefs that contain ResourceVariables.

1) In the ResourceVariable implementation, pass import_scope when creating the SaveSliceInfo. This is present in the implementation of plain variables, and was likely a copy-and-paste omission. 2) When importing a MetaGraphDef, restoring the GLOBAL_VARIABLES and TRAINABLE_VARIABLES collections will add ops to the graph for ResourceVariables. Made graph construction deterministic by fixing the order in which collections are restored. PiperOrigin-RevId: 177144138

Author: tensorflower-gardener

tensorflow/python/framework/meta_graph.py

@@ -663,7 +663,7 @@ def import_scoped_meta_graph(meta_graph_or_file,
  [part for part in [graph.get_name_scope(), import_scope] if part])
 
  # Restores all the other collections.
- for key, col_def in meta_graph_def.collection_def.items():
+ for key, col_def in sorted(meta_graph_def.collection_def.items()):
  # Don't add unbound_inputs to the new graph.
  if key == unbound_inputs_col_name:
  continue

tensorflow/python/framework/meta_graph_test.py

@@ -662,22 +662,36 @@ def _enqueue_vector(sess, queue, values, shape=None):
 class ExportImportAcrossScopesTest(test.TestCase):
 
  def testPartionedVariables(self):
- def make_graph_with_partitioned_variables():
+
+ def make_graph_with_partitioned_variables(use_resource):
  variable_scope.get_variable(
  name="weights",
  partitioner=partitioned_variables.fixed_size_partitioner(3, axis=0),
- initializer=random_ops.truncated_normal([100, 10]))
- self._testExportImportAcrossScopes(make_graph_with_partitioned_variables)
+ initializer=random_ops.truncated_normal([100, 10]),
+ use_resource=use_resource)
+ # The next variable illustrates the necessity of restoring collections
+ # in a deterministic fashion when using ResourceVariables.
+ variable_scope.get_variable(
+ name="another",
+ shape=[],
+ collections=["a", "b", "z", "f", "e", "d", "g"],
+ use_resource=use_resource)
+
+ self._testExportImportAcrossScopes(
+ make_graph_with_partitioned_variables, use_resource=False)
+ self._testExportImportAcrossScopes(
+ make_graph_with_partitioned_variables, use_resource=True)
 
- def _testExportImportAcrossScopes(self, graph_fn):
+ def _testExportImportAcrossScopes(self, graph_fn, use_resource):
  """Tests export and importing a graph across scopes.
 
  Args:
  graph_fn: A closure that creates a graph on the current scope.
+ use_resource: A bool indicating whether or not to use ResourceVariables.
  """
  with ops.Graph().as_default() as original_graph:
  with variable_scope.variable_scope("dropA/dropB/keepA"):
- graph_fn()
+ graph_fn(use_resource=use_resource)
  exported_meta_graph_def = meta_graph.export_scoped_meta_graph(
  graph=original_graph,
  export_scope="dropA/dropB")[0]
@@ -689,10 +703,32 @@ def _testExportImportAcrossScopes(self, graph_fn):
 
  with ops.Graph().as_default() as expected_graph:
  with variable_scope.variable_scope("importA/keepA"):
- graph_fn()
+ graph_fn(use_resource=use_resource)
+
+ if use_resource:
+ # Bringing in a collection that contains ResourceVariables adds ops
+ # to the graph, so mimic the same behavior.
+ for collection_key in sorted([
+ ops.GraphKeys.GLOBAL_VARIABLES,
+ ops.GraphKeys.TRAINABLE_VARIABLES,
+ ]):
+ for var in expected_graph.get_collection(collection_key):
+ var._read_variable_op()
 
  result = meta_graph.export_scoped_meta_graph(graph=imported_graph)[0]
  expected = meta_graph.export_scoped_meta_graph(graph=expected_graph)[0]
+
+ if use_resource:
+ # Clear all shared_name attributes before comparing, since they are
+ # supposed to be orthogonal to scopes.
+ for meta_graph_def in [result, expected]:
+ for node in meta_graph_def.graph_def.node:
+ shared_name_attr = "shared_name"
+ 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 = b""
+
  self.assertProtoEquals(expected, result)
 
 

tensorflow/python/ops/resource_variable_ops.py

@@ -513,7 +513,8 @@ def _init_from_proto(self, variable_def, import_scope=None):
  self._cached_value = None
  if variable_def.HasField("save_slice_info_def"):
  self._save_slice_info = variables.Variable.SaveSliceInfo(
- save_slice_info_def=variable_def.save_slice_info_def)
+ save_slice_info_def=variable_def.save_slice_info_def,
+ import_scope=import_scope)
  else:
  self._save_slice_info = None
  self._caching_device = None