diff options
author | 2018-08-10 15:57:45 -0700 | |
---|---|---|
committer | 2018-08-10 16:01:41 -0700 | |
commit | 8d532ac40f4db7f5293610fd3c6e92a3f7409b76 (patch) | |
tree | f0a57897cac3baa4259ff8a9293befec3dcf1d35 /tensorflow/core/kernels/data/flat_map_dataset_op.cc | |
parent | 84af5e7061f82240828f72c7b484a1a66b8c4f7f (diff) |
[tf.data] Optimization checkpointing improvements.
This CL:
- changes the `OptimizeDataset` checkpointing logic to checkpoint the optimized dataset (as opposed to the original dataset + the optimizations, re-running optimization every time a checkpoint is restored)
- replaces `OpKernelContext` with newly introduced `SerializationContext` in the signature of `AsGraphDefInternal` to reduce the scope of the context and also simplify the logic for overriding the `FunctionLibraryDefinition` when optimizations take place
PiperOrigin-RevId: 208282562
Diffstat (limited to 'tensorflow/core/kernels/data/flat_map_dataset_op.cc')
-rw-r--r-- | tensorflow/core/kernels/data/flat_map_dataset_op.cc | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/tensorflow/core/kernels/data/flat_map_dataset_op.cc b/tensorflow/core/kernels/data/flat_map_dataset_op.cc index aae3f19c0d..3419eed6c6 100644 --- a/tensorflow/core/kernels/data/flat_map_dataset_op.cc +++ b/tensorflow/core/kernels/data/flat_map_dataset_op.cc @@ -91,9 +91,10 @@ class FlatMapDatasetOp : public UnaryDatasetOpKernel { string DebugString() const override { return "FlatMapDatasetOp::Dataset"; } protected: - Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b, + Status AsGraphDefInternal(SerializationContext* ctx, + DatasetGraphDefBuilder* b, Node** output) const override { - TF_RETURN_IF_ERROR(b->AddFunction(ctx, func_.name())); + TF_RETURN_IF_ERROR(b->AddFunction(ctx->flib_def(), func_.name())); Node* input_graph_node = nullptr; TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph_node)); |