Keeps the status of a name-based checkpoint restore.
| 296 | |
| 297 | |
| 298 | class _NameBasedRestoreCoordinator(object): |
| 299 | """Keeps the status of a name-based checkpoint restore.""" |
| 300 | |
| 301 | def __init__(self, save_path, dtype_map=None): |
| 302 | self.save_path = save_path |
| 303 | self.dtype_map = dtype_map |
| 304 | # A map from trackable objects to unused attribute names. We don't have |
| 305 | # proto IDs when doing a name-based restore, so the map keys differ from |
| 306 | # those in _CheckpointRestoreCoordinator. |
| 307 | self.unused_attributes = object_identity.ObjectIdentityWeakKeyDictionary() |
| 308 | self.restore_uid = ops.uid() |
| 309 | |
| 310 | def globally_named_object_attributes(self, trackable): |
| 311 | """Create globally named SaveableObjects from attributes. |
| 312 | |
| 313 | If an object's attribute has no global name specified (default construction |
| 314 | for the SaveableObject factory), records the failure in |
| 315 | `self.unused_attributes` (which can then be used to make status assertions |
| 316 | fail; see `NameBasedSaverStatus`). |
| 317 | |
| 318 | Args: |
| 319 | trackable: An object to save. |
| 320 | |
| 321 | Yields: |
| 322 | SaveableObjects for `trackable`'s attributes. |
| 323 | """ |
| 324 | for attribute_name, saveable_factory in ( |
| 325 | trackable._gather_saveables_for_checkpoint().items()): # pylint: disable=protected-access |
| 326 | if callable(saveable_factory): |
| 327 | try: |
| 328 | # This saveable object factory does not have a default name= argument, |
| 329 | # which means there's no way to save/restore it using a name-based |
| 330 | # checkpoint. Ignore the error now and make sure assert_consumed() |
| 331 | # fails. |
| 332 | saveable = saveable_factory() |
| 333 | except TypeError: |
| 334 | # Even if we can't name this object, we should construct it and check |
| 335 | # whether it's optional to restore it. If it's optional we don't need |
| 336 | # to make assertions fail. |
| 337 | if not saveable_factory("").optional_restore: |
| 338 | self.unused_attributes.setdefault(trackable, |
| 339 | []).append(attribute_name) |
| 340 | continue |
| 341 | else: |
| 342 | saveable = saveable_factory |
| 343 | names_to_saveables = saveable_object_util.op_list_to_dict( |
| 344 | [saveable], convert_variable_to_tensor=False) |
| 345 | for name, op in names_to_saveables.items(): |
| 346 | for saveable_object in saveable_object_util.saveable_objects_for_op( |
| 347 | op=op, name=name): |
| 348 | yield saveable_object |
| 349 | |
| 350 | def eager_restore(self, trackable): |
| 351 | """Runs restore ops for `trackable`'s attributes.""" |
| 352 | # When graph building, we don't add any restore ops to the graph until |
| 353 | # run_restore_ops/initialize_or_restore on the status object for name-based |
| 354 | # checkpoints. |
| 355 | assert context.executing_eagerly() |