Load state from checkpoint into the deserialized objects.
(self)
| 270 | setattr(type(obj), "__call__", _call_attribute) |
| 271 | |
| 272 | def _restore_checkpoint(self): |
| 273 | """Load state from checkpoint into the deserialized objects.""" |
| 274 | variables_path = saved_model_utils.get_variables_path(self._export_dir) |
| 275 | # TODO(andresp): Clean use of private methods of TrackableSaver. |
| 276 | # pylint: disable=protected-access |
| 277 | saver = util.TrackableSaver(graph_view.ObjectGraphView(self.get(0))) |
| 278 | with ops.device("CPU"): |
| 279 | saver._file_prefix_placeholder = constant_op.constant(variables_path) |
| 280 | load_status = saver.restore(variables_path) |
| 281 | load_status.assert_existing_objects_matched() |
| 282 | checkpoint = load_status._checkpoint |
| 283 | |
| 284 | # When running in eager mode, the `restore` call above has already run and |
| 285 | # restored the state of trackables, call `position.restore_ops()` will |
| 286 | # return an empty list as there is nothing left to do. In graph mode, that |
| 287 | # will return the list of ops that must run to restore the object on that |
| 288 | # position. We have to wire them in the initializers of the objects so that |
| 289 | # they get initialized properly when using common practices (e.g. the ones |
| 290 | # used by ManagedSession) without further user action. |
| 291 | for object_id, obj in dict(checkpoint.object_by_proto_id).items(): |
| 292 | position = base.CheckpointPosition(checkpoint=checkpoint, |
| 293 | proto_id=object_id) |
| 294 | restore_ops = position.restore_ops() |
| 295 | if restore_ops: |
| 296 | if resource_variable_ops.is_resource_variable(obj): |
| 297 | obj._initializer_op = restore_ops |
| 298 | else: |
| 299 | raise NotImplementedError( |
| 300 | ("Missing functionality to restore state of object " |
| 301 | "%r from the checkpoint." % obj)) |
| 302 | |
| 303 | def get(self, node_id): |
| 304 | return self._nodes[node_id] |
no test coverage detected