EVOLUTION-MANAGER

Edit File: graph_view.py

"""Manages a graph of Trackable objects.""" # 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. # ============================================================================== from __future__ import absolute_import from __future__ import division from __future__ import print_function import collections import weakref from tensorflow.core.protobuf import trackable_object_graph_pb2 from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.training import optimizer as optimizer_v1 from tensorflow.python.training.saving import saveable_object as saveable_object_lib from tensorflow.python.training.saving import saveable_object_util from tensorflow.python.training.tracking import base from tensorflow.python.training.tracking import tracking from tensorflow.python.util import object_identity _ESCAPE_CHAR = "." # For avoiding conflicts with user-specified names. # Keyword for identifying that the next bit of a checkpoint variable name is a # slot name. Checkpoint names for slot variables look like: # # <path to variable>/<_OPTIMIZER_SLOTS_NAME>/<path to optimizer>/<slot name> # # Where <path to variable> is a full path from the checkpoint root to the # variable being slotted for. _OPTIMIZER_SLOTS_NAME = _ESCAPE_CHAR + "OPTIMIZER_SLOT" # Keyword for separating the path to an object from the name of an # attribute in checkpoint names. Used like: # <path to variable>/<_OBJECT_ATTRIBUTES_NAME>/<name of attribute> _OBJECT_ATTRIBUTES_NAME = _ESCAPE_CHAR + "ATTRIBUTES" def _escape_local_name(name): # We need to support slashes in local names for compatibility, since this # naming scheme is being patched in to things like Layer.add_variable where # slashes were previously accepted. We also want to use slashes to indicate # edges traversed to reach the variable, so we escape forward slashes in # names. return (name.replace(_ESCAPE_CHAR, _ESCAPE_CHAR + _ESCAPE_CHAR) .replace(r"/", _ESCAPE_CHAR + "S")) def _object_prefix_from_path(path_to_root): return "/".join( (_escape_local_name(trackable.name) for trackable in path_to_root)) def _slot_variable_naming_for_optimizer(optimizer_path): """Make a function for naming slot variables in an optimizer.""" # Name slot variables: # # <variable name>/<_OPTIMIZER_SLOTS_NAME>/<optimizer path>/<slot name> # # where <variable name> is exactly the checkpoint name used for the original # variable, including the path from the checkpoint root and the local name in # the object which owns it. Note that we only save slot variables if the # variable it's slotting for is also being saved. optimizer_identifier = "/%s/%s/" % (_OPTIMIZER_SLOTS_NAME, optimizer_path) def _name_slot_variable(variable_path, slot_name): """With an optimizer specified, name a slot variable.""" return (variable_path + optimizer_identifier + _escape_local_name(slot_name)) return _name_slot_variable def _serialize_slot_variables(trackable_objects, node_ids, object_names): """Gather and name slot variables.""" non_slot_objects = list(trackable_objects) slot_variables = object_identity.ObjectIdentityDictionary() for trackable in non_slot_objects: if (isinstance(trackable, optimizer_v1.Optimizer) # TODO(b/110718070): Fix Keras imports. # Note: dir() is used rather than hasattr() here to avoid triggering # custom __getattr__ code, see b/152031870 for context. or "_create_or_restore_slot_variable" in dir(trackable)): naming_scheme = _slot_variable_naming_for_optimizer( optimizer_path=object_names[trackable]) slot_names = trackable.get_slot_names() for slot_name in slot_names: for original_variable_node_id, original_variable in enumerate( non_slot_objects): try: slot_variable = trackable.get_slot( original_variable, slot_name) except (AttributeError, KeyError): slot_variable = None if slot_variable is None: continue slot_variable._maybe_initialize_trackable() # pylint: disable=protected-access if slot_variable._checkpoint_dependencies: # pylint: disable=protected-access # TODO(allenl): Gather dependencies of slot variables. raise NotImplementedError( "Currently only variables with no dependencies can be saved as " "slot variables. File a feature request if this limitation " "bothers you.") if slot_variable in node_ids: raise NotImplementedError( "A slot variable was re-used as a dependency of a " "Trackable object. This is not currently allowed. File a " "feature request if this limitation bothers you.") checkpoint_name = naming_scheme( variable_path=object_names[original_variable], slot_name=slot_name) object_names[slot_variable] = checkpoint_name slot_variable_node_id = len(trackable_objects) node_ids[slot_variable] = slot_variable_node_id trackable_objects.append(slot_variable) slot_variable_proto = ( trackable_object_graph_pb2.TrackableObjectGraph .TrackableObject.SlotVariableReference( slot_name=slot_name, original_variable_node_id=original_variable_node_id, slot_variable_node_id=slot_variable_node_id)) slot_variables.setdefault(trackable, []).append( slot_variable_proto) return slot_variables class ObjectGraphView(object): """Gathers and serializes an object graph.""" def __init__(self, root, saveables_cache=None, attached_dependencies=None): """Configure the graph view. Args: root: A `Trackable` object whose variables (including the variables of dependencies, recursively) should be saved. May be a weak reference. saveables_cache: A dictionary mapping `Trackable` objects -> attribute names -> SaveableObjects, used to avoid re-creating SaveableObjects when graph building. attached_dependencies: Dependencies to attach to the root object. Used when saving a Checkpoint with a defined root object. """ self._root_ref = root self._saveables_cache = saveables_cache self._attached_dependencies = attached_dependencies def list_dependencies(self, obj): # pylint: disable=protected-access obj._maybe_initialize_trackable() dependencies = obj._checkpoint_dependencies # pylint: enable=protected-access if obj is self.root and self._attached_dependencies: dependencies = dependencies.copy() dependencies.extend(self._attached_dependencies) return dependencies @property def saveables_cache(self): """Maps Trackable objects -> attribute names -> list(SaveableObjects). Used to avoid re-creating SaveableObjects when graph building. None when executing eagerly. Returns: The cache (an object-identity dictionary), or None if caching is disabled. """ return self._saveables_cache @property def attached_dependencies(self): """Returns list of dependencies that should be saved in the checkpoint. These dependencies are not tracked by root, but are in the the checkpoint. This is defined when the user creates a Checkpoint with both root and kwargs set. Returns: A list of TrackableReferences. """ return self._attached_dependencies @property def root(self): if isinstance(self._root_ref, weakref.ref): derefed = self._root_ref() assert derefed is not None return derefed else: return self._root_ref def _breadth_first_traversal(self): """Find shortest paths to all dependencies of self.root.""" bfs_sorted = [] to_visit = collections.deque([self.root]) path_to_root = object_identity.ObjectIdentityDictionary() path_to_root[self.root] = () while to_visit: current_trackable = to_visit.popleft() if isinstance(current_trackable, tracking.NotTrackable): raise NotImplementedError( ("The object %s does not support object-based saving. File a " "feature request if this limitation bothers you. In the meantime, " "you can remove the dependency on this object and save everything " "else.") % (current_trackable,)) bfs_sorted.append(current_trackable) for name, dependency in self.list_dependencies(current_trackable): if dependency not in path_to_root: path_to_root[dependency] = ( path_to_root[current_trackable] + ( base.TrackableReference(name, dependency),)) to_visit.append(dependency) return bfs_sorted, path_to_root def _add_attributes_to_object_graph( self, trackable_objects, object_graph_proto, node_ids, object_names, object_map, call_with_mapped_captures): """Create SaveableObjects and corresponding SerializedTensor protos.""" named_saveable_objects = [] if self._saveables_cache is None: # No SaveableObject caching. Either we're executing eagerly, or building a # static save which is specialized to the current Python state. feed_additions = None else: # If we are caching SaveableObjects, we need to build up a feed_dict with # functions computing volatile Python state to be saved with the # checkpoint. feed_additions = {} for checkpoint_id, (trackable, object_proto) in enumerate( zip(trackable_objects, object_graph_proto.nodes)): assert node_ids[trackable] == checkpoint_id object_name = object_names[trackable] if object_map is None: object_to_save = trackable else: object_to_save = object_map.get(trackable, trackable) if self._saveables_cache is not None: cached_attributes = self._saveables_cache.setdefault(object_to_save, {}) else: cached_attributes = None for name, saveable_factory in ( object_to_save._gather_saveables_for_checkpoint().items()): # pylint: disable=protected-access attribute = object_proto.attributes.add() attribute.name = name attribute.checkpoint_key = "%s/%s/%s" % ( object_name, _OBJECT_ATTRIBUTES_NAME, _escape_local_name(name)) if cached_attributes is None: saveables = None else: saveables = cached_attributes.get(name, None) if saveables is not None: for saveable in saveables: if attribute.checkpoint_key not in saveable.name: # The checkpoint key for this SaveableObject is different. We # need to re-create it. saveables = None del cached_attributes[name] break if saveables is None: if callable(saveable_factory): maybe_saveable = saveable_object_util.create_saveable_object( saveable_factory, attribute.checkpoint_key, call_with_mapped_captures) else: maybe_saveable = saveable_factory if isinstance(maybe_saveable, saveable_object_lib.SaveableObject): saveables = (maybe_saveable,) else: # Figure out the name-based Saver's name for this variable. If it's # already a SaveableObject we'd just get the checkpoint key back, so # we leave full_name blank. saver_dict = saveable_object_util.op_list_to_dict( [maybe_saveable], convert_variable_to_tensor=False) full_name, = saver_dict.keys() saveables = tuple(saveable_object_util.saveable_objects_for_op( op=maybe_saveable, name=attribute.checkpoint_key)) for saveable in saveables: saveable.full_name = full_name for saveable in saveables: if attribute.checkpoint_key not in saveable.name: raise AssertionError( ("The object %s produced a SaveableObject with name '%s' for " "attribute '%s'. Expected a name containing '%s'.") % (trackable, name, saveable.name, attribute.checkpoint_key)) if cached_attributes is not None: cached_attributes[name] = saveables optional_restore = None for saveable in saveables: if optional_restore is None: optional_restore = saveable.optional_restore else: optional_restore = optional_restore and saveable.optional_restore if hasattr(saveable, "full_name"): attribute.full_name = saveable.full_name if isinstance(saveable, base.PythonStateSaveable): if feed_additions is None: assert self._saveables_cache is None # If we're not caching saveables, then we're either executing # eagerly or building a static save/restore (e.g. for a # SavedModel). In either case, we should embed the current Python # state in the graph rather than relying on a feed dict. saveable = saveable.freeze() else: saveable_feed_dict = saveable.feed_dict_additions() for new_feed_key in saveable_feed_dict.keys(): if new_feed_key in feed_additions: raise AssertionError( ("The object %s tried to feed a value for the Tensor %s " "when saving, but another object is already feeding a " "value.") % (trackable, new_feed_key)) feed_additions.update(saveable_feed_dict) named_saveable_objects.append(saveable) if optional_restore is None: optional_restore = False attribute.optional_restore = optional_restore return named_saveable_objects, feed_additions def _fill_object_graph_proto(self, trackable_objects, node_ids, slot_variables, object_graph_proto=None): """Name non-slot `Trackable`s and add them to `object_graph_proto`.""" if object_graph_proto is None: object_graph_proto = ( trackable_object_graph_pb2.TrackableObjectGraph()) for checkpoint_id, trackable in enumerate(trackable_objects): assert node_ids[trackable] == checkpoint_id object_proto = object_graph_proto.nodes.add() object_proto.slot_variables.extend(slot_variables.get(trackable, ())) for child in self.list_dependencies(trackable): child_proto = object_proto.children.add() child_proto.node_id = node_ids[child.ref] child_proto.local_name = child.name return object_graph_proto def _serialize_gathered_objects(self, trackable_objects, path_to_root, object_map=None, call_with_mapped_captures=None): """Create SaveableObjects and protos for gathered objects.""" object_names = object_identity.ObjectIdentityDictionary() for obj, path in path_to_root.items(): object_names[obj] = _object_prefix_from_path(path) node_ids = object_identity.ObjectIdentityDictionary() for node_id, node in enumerate(trackable_objects): node_ids[node] = node_id slot_variables = _serialize_slot_variables( trackable_objects=trackable_objects, node_ids=node_ids, object_names=object_names) object_graph_proto = self._fill_object_graph_proto( trackable_objects=trackable_objects, node_ids=node_ids, slot_variables=slot_variables) named_saveable_objects, feed_additions = ( self._add_attributes_to_object_graph( trackable_objects=trackable_objects, object_graph_proto=object_graph_proto, node_ids=node_ids, object_names=object_names, object_map=object_map, call_with_mapped_captures=call_with_mapped_captures)) return named_saveable_objects, object_graph_proto, feed_additions def serialize_object_graph(self): """Determine checkpoint keys for variables and build a serialized graph. Non-slot variables are keyed based on a shortest path from the root saveable to the object which owns the variable (i.e. the one which called `Trackable._add_variable` to create it). Slot variables are keyed based on a shortest path to the variable being slotted for, a shortest path to their optimizer, and the slot name. Returns: A tuple of (named_variables, object_graph_proto, feed_additions): named_variables: A dictionary mapping names to variable objects. object_graph_proto: A TrackableObjectGraph protocol buffer containing the serialized object graph and variable references. feed_additions: A dictionary mapping from Tensors to values which should be fed when saving. Raises: ValueError: If there are invalid characters in an optimizer's slot names. """ trackable_objects, path_to_root = self._breadth_first_traversal() return self._serialize_gathered_objects( trackable_objects, path_to_root) def frozen_saveable_objects(self, object_map=None, to_graph=None, call_with_mapped_captures=None): """Creates SaveableObjects with the current object graph frozen.""" trackable_objects, path_to_root = self._breadth_first_traversal() if to_graph: target_context = to_graph.as_default else: target_context = ops.NullContextmanager with target_context(): named_saveable_objects, graph_proto, _ = self._serialize_gathered_objects( trackable_objects, path_to_root, object_map, call_with_mapped_captures) with ops.device("/cpu:0"): object_graph_tensor = constant_op.constant( graph_proto.SerializeToString(), dtype=dtypes.string) named_saveable_objects.append( base.NoRestoreSaveable( tensor=object_graph_tensor, name=base.OBJECT_GRAPH_PROTO_KEY)) return named_saveable_objects def objects_ids_and_slot_variables(self): """Traverse the object graph and list all accessible objects. Looks for `Trackable` objects which are dependencies of `root_trackable`. Includes slot variables only if the variable they are slotting for and the optimizer are dependencies of `root_trackable` (i.e. if they would be saved with a checkpoint). Returns: A tuple of (trackable objects, object -> node id, slot variables) """ trackable_objects, path_to_root = self._breadth_first_traversal() object_names = object_identity.ObjectIdentityDictionary() for obj, path in path_to_root.items(): object_names[obj] = _object_prefix_from_path(path) node_ids = object_identity.ObjectIdentityDictionary() for node_id, node in enumerate(trackable_objects): node_ids[node] = node_id slot_variables = _serialize_slot_variables( trackable_objects=trackable_objects, node_ids=node_ids, object_names=object_names) return trackable_objects, node_ids, slot_variables def list_objects(self): """Traverse the object graph and list all accessible objects.""" trackable_objects, _, _ = self.objects_ids_and_slot_variables() return trackable_objects