Katherine Wu a3e64f721c Add SaveableObjects to SavedModel. ...

When objects are loaded from the SavedModel, they don't retain their `_gather_saveables_for_checkpoint` functions, which can result in values not being loaded from the checkpoint.

This CL adds a field in the SavedModel proto that stores a save and restore function for each SaveableObject in each node. When loading into Python, the SaveableObjects are restored using the functions.

PiperOrigin-RevId: 318549786
Change-Id: I688c72d7658e1bca98abf373a13a0e15a7fb83e2

2020-07-01 14:28:41 -07:00

..

BUILD

Wrap save/restore logic in tf.function when in eager mode. This allows parallel saving and restoring when using multiple devices.

2020-06-22 13:23:14 -07:00

checkpoint_options.py

Add an option to choose the I/O Device for saving and loading models.

2020-04-22 11:27:15 -07:00

functional_saver_test.py

Wrap save/restore logic in tf.function when in eager mode. This allows parallel saving and restoring when using multiple devices.

2020-06-22 13:23:14 -07:00

functional_saver.py

Wrap save/restore logic in tf.function when in eager mode. This allows parallel saving and restoring when using multiple devices.

2020-06-22 13:23:14 -07:00

saveable_hook.py

Add SaveableHook, a special SaveableObject which registers callbacks.

2019-11-15 12:18:18 -08:00

saveable_object_util.py

Add SaveableObjects to SavedModel.

2020-07-01 14:28:41 -07:00

saveable_object.py

In all cases, we can't rely on the tensor.device attribute being set. So its better to get the device for a SaveSpec from the device passed in rather. This was an issue with saving iterators because for iterators the resource usually has a device specification but the serialized tensor derived from it might not have it set. As a result, when saving iterators in a sharded fashion all iterators end up on '' device instead which is not what is intended.

2019-12-18 19:11:04 -08:00