Provides a frozen view over a trackable root. This class helps creating a single stable view over an object to save. The saving code should access properties and functions via this class and not via the original object as there are cases where an object construct their trackable attributes
| 143 | |
| 144 | |
| 145 | class _SaveableView(object): |
| 146 | """Provides a frozen view over a trackable root. |
| 147 | |
| 148 | This class helps creating a single stable view over an object to save. The |
| 149 | saving code should access properties and functions via this class and not via |
| 150 | the original object as there are cases where an object construct their |
| 151 | trackable attributes and functions dynamically per call and will yield |
| 152 | different objects if invoked more than once. |
| 153 | |
| 154 | Changes to the graph, for example adding objects, must happen in |
| 155 | `checkpoint_view` (an `_AugmentedGraphView`) before the `_SaveableView` is |
| 156 | constructed. Changes after the `_SaveableView` has been constructed will be |
| 157 | ignored. |
| 158 | """ |
| 159 | |
| 160 | def __init__(self, checkpoint_view): |
| 161 | self.checkpoint_view = checkpoint_view |
| 162 | trackable_objects, node_ids, slot_variables = ( |
| 163 | self.checkpoint_view.objects_ids_and_slot_variables()) |
| 164 | self.nodes = trackable_objects |
| 165 | self.node_ids = node_ids |
| 166 | self.captured_tensor_node_ids = object_identity.ObjectIdentityDictionary() |
| 167 | self.slot_variables = slot_variables |
| 168 | self.concrete_functions = [] |
| 169 | |
| 170 | # Also add `Function`s as nodes. |
| 171 | nodes_without_functions = list(self.nodes) |
| 172 | seen_function_names = set() |
| 173 | for node in nodes_without_functions: |
| 174 | for function in checkpoint_view.list_functions(node).values(): |
| 175 | if function not in self.node_ids: |
| 176 | self.node_ids[function] = len(self.nodes) |
| 177 | self.nodes.append(function) |
| 178 | if isinstance(function, def_function.Function): |
| 179 | # Force listing the concrete functions for the side effects: |
| 180 | # - populate the cache for functions that have an input_signature |
| 181 | # and have not been called. |
| 182 | # - force side effects of creation of concrete functions, e.g. create |
| 183 | # variables on first run. |
| 184 | concrete_functions = ( |
| 185 | function._list_all_concrete_functions_for_serialization()) # pylint: disable=protected-access |
| 186 | else: |
| 187 | concrete_functions = [function] |
| 188 | for concrete_function in concrete_functions: |
| 189 | if concrete_function.name not in seen_function_names: |
| 190 | seen_function_names.add(concrete_function.name) |
| 191 | self.concrete_functions.append(concrete_function) |
| 192 | |
| 193 | @property |
| 194 | def root(self): |
| 195 | return self.nodes[0] |
| 196 | |
| 197 | def fill_object_graph_proto(self, proto): |
| 198 | """Populate the nodes, children and slot_variables of a SavedObjectGraph.""" |
| 199 | for node_id, node in enumerate(self.nodes): |
| 200 | assert self.node_ids[node] == node_id |
| 201 | object_proto = proto.nodes.add() |
| 202 | object_proto.slot_variables.extend(self.slot_variables.get(node, ())) |