An extendable graph which also tracks functions attached to objects. Extensions through `add_object` appear in the object graph and any checkpoints generated from it, even if they are not dependencies of the node they were attached to in the saving program. For example a `.signatures` attribu
| 68 | |
| 69 | |
| 70 | class _AugmentedGraphView(graph_view.ObjectGraphView): |
| 71 | """An extendable graph which also tracks functions attached to objects. |
| 72 | |
| 73 | Extensions through `add_object` appear in the object graph and any checkpoints |
| 74 | generated from it, even if they are not dependencies of the node they were |
| 75 | attached to in the saving program. For example a `.signatures` attribute is |
| 76 | added to exported SavedModel root objects without modifying the root object |
| 77 | itself. |
| 78 | |
| 79 | Also tracks functions attached to objects in the graph, through the caching |
| 80 | `list_functions` method. Enumerating functions only through this method |
| 81 | ensures that we get a consistent view of functions, even if object attributes |
| 82 | create new functions every time they are accessed. |
| 83 | """ |
| 84 | |
| 85 | def __init__(self, root): |
| 86 | if (not context.executing_eagerly() |
| 87 | and not ops.inside_function()): |
| 88 | saveables_cache = object_identity.ObjectIdentityWeakKeyDictionary() |
| 89 | else: |
| 90 | saveables_cache = None |
| 91 | super(_AugmentedGraphView, self).__init__(root, saveables_cache) |
| 92 | # Object -> (name -> dep) |
| 93 | self._extra_dependencies = object_identity.ObjectIdentityDictionary() |
| 94 | self._functions = object_identity.ObjectIdentityDictionary() |
| 95 | # Cache shared between objects in the same object graph. This is passed to |
| 96 | # each trackable object's `_list_extra_dependencies_for_serialization` and |
| 97 | # `_list_functions_for_serialization` function. |
| 98 | self._serialization_cache = object_identity.ObjectIdentityDictionary() |
| 99 | |
| 100 | def add_object(self, parent_node, name_in_parent, subgraph_root): |
| 101 | """Attach an object to `parent_node`, overriding any existing dependency.""" |
| 102 | self._extra_dependencies.setdefault( |
| 103 | parent_node, {})[name_in_parent] = subgraph_root |
| 104 | |
| 105 | def list_dependencies(self, obj): |
| 106 | """Overrides a parent method to include `add_object` objects.""" |
| 107 | extra_dependencies = self.list_extra_dependencies(obj) |
| 108 | extra_dependencies.update(self._extra_dependencies.get(obj, {})) |
| 109 | |
| 110 | used_names = set() |
| 111 | for name, dep in super(_AugmentedGraphView, self).list_dependencies(obj): |
| 112 | used_names.add(name) |
| 113 | if name in extra_dependencies: |
| 114 | # Extra dependencies (except for `.signatures`, which is always added |
| 115 | # when saving) should not have naming conflicts with dependencies |
| 116 | # defined by the user. |
| 117 | if name != signature_serialization.SIGNATURE_ATTRIBUTE_NAME: |
| 118 | raise ValueError( |
| 119 | "Error when exporting object {} of with identifier={}. The object" |
| 120 | " has an attribute named {}, which is reserved. List of all " |
| 121 | "reserved attributes: {}".format( |
| 122 | obj, obj._object_identifier, # pylint: disable=protected-access |
| 123 | name, extra_dependencies.keys())) |
| 124 | yield base.TrackableReference(name, extra_dependencies[name]) |
| 125 | else: |
| 126 | yield base.TrackableReference(name, dep) |
| 127 | for name, dep in extra_dependencies.items(): |