Status for loading a name-based training checkpoint.
| 918 | |
| 919 | |
| 920 | class NameBasedSaverStatus(_LoadStatus): |
| 921 | """Status for loading a name-based training checkpoint.""" |
| 922 | |
| 923 | # Ideally this deprecation decorator would be on the class, but that |
| 924 | # interferes with isinstance checks. |
| 925 | @deprecation.deprecated( |
| 926 | date=None, instructions=_DEPRECATED_RESTORE_INSTRUCTIONS) |
| 927 | def __init__(self, checkpoint, graph_view): |
| 928 | self._checkpoint = checkpoint |
| 929 | self._graph_view = graph_view |
| 930 | # Keep a reference to the root, since graph_view might only have a weakref. |
| 931 | self._root = graph_view.root |
| 932 | |
| 933 | |
| 934 | def assert_consumed(self): |
| 935 | """Raises an exception if any variables/objects are unmatched.""" |
| 936 | unused_attributes = list(self._checkpoint.unused_attributes.items()) |
| 937 | if unused_attributes: |
| 938 | unused_attribute_strings = [ |
| 939 | "\n {}: {}".format(obj, attributes) |
| 940 | for obj, attributes in unused_attributes] |
| 941 | raise AssertionError( |
| 942 | "Some objects had attributes which were not restored:{}".format( |
| 943 | "".join(unused_attribute_strings))) |
| 944 | for trackable in self._graph_view.list_objects(): |
| 945 | # pylint: disable=protected-access |
| 946 | trackable._maybe_initialize_trackable() |
| 947 | if trackable._update_uid < self._checkpoint.restore_uid: |
| 948 | raise AssertionError("Object not restored: %s" % (trackable,)) |
| 949 | # pylint: enable=protected-access |
| 950 | return self |
| 951 | |
| 952 | def assert_existing_objects_matched(self): |
| 953 | """Raises an exception if currently created objects are unmatched.""" |
| 954 | # For name-based checkpoints there's no object information in the |
| 955 | # checkpoint, so there's no distinction between |
| 956 | # assert_existing_objects_matched and assert_consumed (and both are less |
| 957 | # useful since we don't touch Python objects or Python state). |
| 958 | return self.assert_consumed() |
| 959 | |
| 960 | def assert_nontrivial_match(self): |
| 961 | """Raises an exception if currently created objects are unmatched.""" |
| 962 | # For name-based checkpoints there's no object information in the |
| 963 | # checkpoint, so there's no distinction between |
| 964 | # assert_nontrivial_match and assert_consumed (and both are less |
| 965 | # useful since we don't touch Python objects or Python state). |
| 966 | return self.assert_consumed() |
| 967 | |
| 968 | def _gather_saveable_objects(self): |
| 969 | """Walk the object graph, using global names for SaveableObjects.""" |
| 970 | objects = self._graph_view.list_objects() |
| 971 | saveable_objects = [] |
| 972 | for trackable in objects: |
| 973 | # pylint: disable=protected-access |
| 974 | trackable._maybe_initialize_trackable() |
| 975 | if trackable._update_uid < self._checkpoint.restore_uid: |
| 976 | trackable._update_uid = self._checkpoint.restore_uid |
| 977 | else: |