Serialize the Dataset graph alongside the Iterator state when checkpointing Iterators.
The Dataset graph is stored as a serialized GraphDef using a pre-defined key. The name of the output node of this GraphDef is stored in a separate key using the BundleWriter. When restoring the checkpoint, the Dataset graph, if available, is deserialized and executed using the GraphRunner to get the Dataset which is then used to construct the Iterator. Also moved BundleReaderWrapper and BundleWriterWrapper out of IteratorBase so they can be more generally used. Added a GraphDatasetBase that will be used as a base class for all Datasets that are used only in ops. PiperOrigin-RevId: 170263870
Loading
Please sign in to comment