diff --git a/tensorflow/python/data/experimental/__init__.py b/tensorflow/python/data/experimental/__init__.py index 8cec75b59928ef437e5611381fcb763b5a227e63..8a1048513af379a73162d6a0629274c30a86f8c0 100644 --- a/tensorflow/python/data/experimental/__init__.py +++ b/tensorflow/python/data/experimental/__init__.py @@ -25,6 +25,7 @@ See [Importing Data](https://tensorflow.org/guide/datasets) for an overview. @@Counter @@CheckpointInputPipelineHook @@CsvDataset +@@OptimizationOptions @@Optional @@RandomDataset @@Reducer @@ -86,10 +87,8 @@ from tensorflow.python.data.experimental.ops.interleave_ops import parallel_inte from tensorflow.python.data.experimental.ops.interleave_ops import sample_from_datasets from tensorflow.python.data.experimental.ops.iterator_ops import CheckpointInputPipelineHook from tensorflow.python.data.experimental.ops.iterator_ops import make_saveable_from_iterator - -# Optimization constant that can be used to enable auto-tuning. from tensorflow.python.data.experimental.ops.optimization import AUTOTUNE - +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.experimental.ops.parsing_ops import parse_example_dataset from tensorflow.python.data.experimental.ops.prefetching_ops import copy_to_device from tensorflow.python.data.experimental.ops.prefetching_ops import prefetch_to_device diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD index 121798ad3e6b55e6ca1c600c5345d3ae2a98a10f..e05f382171836bb624145ba6cd4b4c91488a714e 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD @@ -42,6 +42,7 @@ py_test( "//tensorflow/python:errors", "//tensorflow/python:math_ops", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", "@absl_py//absl/testing:parameterized", @@ -68,6 +69,7 @@ py_test( "//tensorflow/python:math_ops", "//tensorflow/python:random_ops", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", "@absl_py//absl/testing:parameterized", @@ -127,6 +129,7 @@ py_test( "//tensorflow/python:client_testlib", "//tensorflow/python:errors", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", ], @@ -148,6 +151,7 @@ py_test( "//tensorflow/python:errors", "//tensorflow/python:math_ops", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", "@absl_py//absl/testing:parameterized", @@ -167,6 +171,7 @@ py_test( "//tensorflow/python:client_testlib", "//tensorflow/python:errors", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", "@absl_py//absl/testing:parameterized", @@ -192,6 +197,7 @@ py_test( "//tensorflow/python:math_ops", "//tensorflow/python:random_ops", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", "@absl_py//absl/testing:parameterized", @@ -227,6 +233,7 @@ py_test( "//tensorflow/python:parsing_ops", "//tensorflow/python:sparse_tensor", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", @@ -272,6 +279,7 @@ py_test( "//tensorflow/python:errors", "//tensorflow/python:math_ops", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", ], @@ -313,6 +321,7 @@ py_test( "//tensorflow/python:client_testlib", "//tensorflow/python:errors", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/kernel_tests:test_base", "//tensorflow/python/data/ops:dataset_ops", ], diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py index fe2c104e94b77f8f9847d5b591b192ee64c363e8..7371cf31dff33a5de18f3268ecdfc91c6a08b29c 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from absl.testing import parameterized from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op @@ -71,7 +72,8 @@ class FilterFusionTest(test_base.DatasetTestBase, parameterized.TestCase): dataset = dataset.cache() options = dataset_ops.Options() - options.experimental_filter_fusion = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.filter_fusion = True dataset = dataset.with_options(options) expected_output = [] for x in range(5): diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py index e86b19438ef45044a9120354b534b55976f45f32..0aacf8bb07cb5cb67fff6eaa9c4184678b45aa95 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from absl.testing import parameterized from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import context @@ -91,7 +92,8 @@ class HoistRandomUniformTest(test_base.DatasetTestBase, parameterized.TestCase): ["Zip[0]", "Map"] if will_optimize else ["Map"])).map(function) options = dataset_ops.Options() - options.experimental_hoist_random_uniform = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.hoist_random_uniform = True dataset = dataset.with_options(options) self._testDataset(dataset) @@ -107,7 +109,8 @@ class HoistRandomUniformTest(test_base.DatasetTestBase, parameterized.TestCase): dataset = dataset_ops.Dataset.range(5).apply( optimization.assert_next(["Zip[0]", "Map"])).map(random_with_capture) options = dataset_ops.Options() - options.experimental_hoist_random_uniform = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.hoist_random_uniform = True dataset = dataset.with_options(options) self._testDataset(dataset) diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py index 67f3ceeabef764130395eafd08e2d08a46a49cb6..801f664f09c7f2f7008090f356a246ca530ddcd5 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import test_util @@ -32,7 +33,8 @@ class MapAndBatchFusionTest(test_base.DatasetTestBase): optimization.assert_next( ["MapAndBatch"])).map(lambda x: x * x).batch(10) options = dataset_ops.Options() - options.experimental_map_and_batch_fusion = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.map_and_batch_fusion = True dataset = dataset.with_options(options) self.assertDatasetProduces( dataset, expected_output=[[x * x for x in range(10)]]) diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py index a898c3844003bb803b1ed38bc18d95b6ddbae0ee..db8f214fbfca1389af70df55518c885610984031 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from absl.testing import parameterized from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op @@ -83,7 +84,8 @@ class MapAndFilterFusionTest(test_base.DatasetTestBase, parameterized.TestCase): optimization.assert_next( ["Map", "FilterByLastComponent"])).map(function).filter(predicate) options = dataset_ops.Options() - options.experimental_map_and_filter_fusion = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.map_and_filter_fusion = True dataset = dataset.with_options(options) self._testMapAndFilter(dataset, function, predicate) @@ -101,7 +103,8 @@ class MapAndFilterFusionTest(test_base.DatasetTestBase, parameterized.TestCase): optimization.assert_next(["Map", "Filter"])).map(function).filter(predicate) options = dataset_ops.Options() - options.experimental_map_and_filter_fusion = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.map_and_filter_fusion = True dataset = dataset.with_options(options) self._testMapAndFilter(dataset, function, predicate) diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py index 47a1b0896cdfd71612b3c3ce4a20e62a93c677f0..d8d63903749d13b80f662c996ebf5c95f934a0b1 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from absl.testing import parameterized from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import test_util @@ -74,7 +75,8 @@ class MapFusionTest(test_base.DatasetTestBase, parameterized.TestCase): dataset = dataset.cache() options = dataset_ops.Options() - options.experimental_map_fusion = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.map_fusion = True dataset = dataset.with_options(options) expected_output = [] for x in range(5): diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py index 042b9ce54bb3ed459691163c722bbfc4565d60ae..0ff3fff4f8550a4221e54ab2b01ddcaf6c340145 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from absl.testing import parameterized from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes @@ -67,7 +68,8 @@ class MapParallelizationTest(test_base.DatasetTestBase, parameterized.TestCase): dataset = dataset_ops.Dataset.range(5).apply( optimization.assert_next(next_nodes)).map(function) options = dataset_ops.Options() - options.experimental_map_parallelization = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.map_parallelization = True dataset = dataset.with_options(options) if should_optimize: self.assertDatasetProduces( diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py index 4f05f026690980fb8985bda7e42d459d51fc193d..d979aaa5a068754a1f685e6fadc6b9d5a67fe5f5 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py @@ -23,6 +23,7 @@ import numpy as np from tensorflow.core.example import example_pb2 from tensorflow.core.example import feature_pb2 from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op @@ -353,7 +354,8 @@ class MapVectorizationTest(test_base.DatasetTestBase, parameterized.TestCase): optimized = _make_dataset(["Batch", map_node_name] if expect_optimized else [map_node_name, "Batch"]) options = dataset_ops.Options() - options.experimental_map_vectorization = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.map_vectorization = True optimized = optimized.with_options(options) return unoptimized, optimized diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py index d957e8007cb4a13b25798b9cb8b4b9e10f910fe7..ce86bfa4e0f8f953722cbb772705ae866ef33e0e 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op @@ -42,7 +43,8 @@ class NoopEliminationTest(test_base.DatasetTestBase): dataset = dataset.repeat(some_tensor).skip(5).take(-1).skip(0).repeat( 1).prefetch(0).prefetch(1).cache() options = dataset_ops.Options() - options.experimental_noop_elimination = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.noop_elimination = True dataset = dataset.with_options(options) self.assertDatasetProduces(dataset, expected_output=range(5)) diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py index f1d00a59c4f8f0b5aa28d789f2b2db1aad53200f..5f746ec63ac8d68d614044e809e7f31178ea8874 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import errors @@ -32,7 +33,8 @@ class ShuffleAndRepeatFusionTest(test_base.DatasetTestBase): dataset = dataset_ops.Dataset.range(10).apply( optimization.assert_next(["ShuffleAndRepeat"])).shuffle(10).repeat(2) options = dataset_ops.Options() - options.experimental_shuffle_and_repeat_fusion = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.shuffle_and_repeat_fusion = True dataset = dataset.with_options(options) get_next = self.getNext(dataset) diff --git a/tensorflow/python/data/experimental/ops/BUILD b/tensorflow/python/data/experimental/ops/BUILD index f85e774887d873f56f7cd6669ef9a631492b8648..50f5127833e607c9fde87f1812019f6f6869e7b3 100644 --- a/tensorflow/python/data/experimental/ops/BUILD +++ b/tensorflow/python/data/experimental/ops/BUILD @@ -235,6 +235,16 @@ py_library( ], ) +py_library( + name = "optimization_options", + srcs = ["optimization_options.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:util", + "//tensorflow/python/data/util:options", + ], +) + py_library( name = "parsing_ops", srcs = ["parsing_ops.py"], diff --git a/tensorflow/python/data/experimental/ops/optimization_options.py b/tensorflow/python/data/experimental/ops/optimization_options.py new file mode 100644 index 0000000000000000000000000000000000000000..dc9d3193748deb8957bdd9a5d0b25f226e6f1955 --- /dev/null +++ b/tensorflow/python/data/experimental/ops/optimization_options.py @@ -0,0 +1,83 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Experimental API for controlling optimizations in `tf.data` pipelines.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + + +from tensorflow.python.data.util import options +from tensorflow.python.util.tf_export import tf_export + + +@tf_export("data.experimental.OptimizationOptions") +class OptimizationOptions(options.OptionsBase): + """Represents options for dataset optimizations. + + You can apply `OptimizationOptions` to a `dataset` object, as follows: + + ```python + options = tf.data.Options() + options.optimization = tf.data.experimental.OptimizationOptions() + options.optimization.map_and_batch_fusion = True + dataset = dataset.with_options(options) + ``` + """ + + filter_fusion = options.create_option( + name="filter_fusion", + ty=bool, + docstring="Whether to fuse filter transformations.") + + hoist_random_uniform = options.create_option( + name="hoist_random_uniform", + ty=bool, + docstring= + "Whether to hoist `tf.random_uniform()` ops out of map transformations.") + + map_and_batch_fusion = options.create_option( + name="map_and_batch_fusion", + ty=bool, + docstring="Whether to fuse map and batch transformations.") + + map_and_filter_fusion = options.create_option( + name="map_and_filter_fusion", + ty=bool, + docstring="Whether to fuse map and filter transformations.") + + map_fusion = options.create_option( + name="map_and_filter_fusion", + ty=bool, + docstring="Whether to fuse map transformations.") + + map_parallelization = options.create_option( + name="map_parallelization", + ty=bool, + docstring="Whether to parallelize stateless map transformations.") + + map_vectorization = options.create_option( + name="map_vectorization", + ty=bool, + docstring="Whether to vectorize map transformations.") + + noop_elimination = options.create_option( + name="noop_elimination", + ty=bool, + docstring="Whether to eliminate no-op transformations.") + + shuffle_and_repeat_fusion = options.create_option( + name="shuffle_and_repeat_fusion", + ty=bool, + docstring="Whether to fuse shuffle and repeat transformations.") diff --git a/tensorflow/python/data/experimental/ops/stats_options.py b/tensorflow/python/data/experimental/ops/stats_options.py index cd7fdcb72343f1e6dc7a6c938fb2ab86c5f4d5b4..6df608c6080d51da7b9a11d0eae7ae0b1d222565 100644 --- a/tensorflow/python/data/experimental/ops/stats_options.py +++ b/tensorflow/python/data/experimental/ops/stats_options.py @@ -34,7 +34,7 @@ class StatsOptions(options.OptionsBase): ```python aggretator = tf.data.experimental.StatsAggregator() - options = dataset_ops.Options() + options = tf.data.Options() options.experimental_stats = tf.data.experimental.StatsOptions() options.experimental_stats.aggregator = aggregator dataset = dataset.with_options(options) diff --git a/tensorflow/python/data/experimental/ops/threading_options.py b/tensorflow/python/data/experimental/ops/threading_options.py index 98df371c2538ae87ab6c3a0f4d515aa3b3ad5b63..dbf662186f818a24a3b19ea678f87351ab45ed6e 100644 --- a/tensorflow/python/data/experimental/ops/threading_options.py +++ b/tensorflow/python/data/experimental/ops/threading_options.py @@ -29,7 +29,7 @@ class ThreadingOptions(options.OptionsBase): To apply `ThreadingOptions` to a `dataset` object, use the following pattern: ```python - options = dataset_ops.Options() + options = tf.data.Options() options.experimental_threading = tf.data.experimental.ThreadingOptions() options.experimental_threading.private_threadpool_size = 10 dataset = dataset.with_options(options) diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD index 0867471d740874d3c562fe1fc73f52d47af6875e..9f7ce99cbc6414480931c92d6bd1f7ead3ec3fe4 100644 --- a/tensorflow/python/data/kernel_tests/BUILD +++ b/tensorflow/python/data/kernel_tests/BUILD @@ -408,6 +408,7 @@ cuda_py_test( "//tensorflow/python/data/ops:multi_device_iterator_ops", "//tensorflow/python/data/ops:iterator_ops", "//tensorflow/python/data/experimental/ops:optimization", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", diff --git a/tensorflow/python/data/kernel_tests/dataset_test.py b/tensorflow/python/data/kernel_tests/dataset_test.py index 7dbab60f9c2de6122977f1e0718753158bad5a77..2952c08be02b76fb221ee0f31f4b9fc34a14d659 100644 --- a/tensorflow/python/data/kernel_tests/dataset_test.py +++ b/tensorflow/python/data/kernel_tests/dataset_test.py @@ -227,12 +227,12 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): options1 = dataset_ops.Options() options1.experimental_autotune = True options2 = dataset_ops.Options() - options2.experimental_filter_fusion = False + options2.experimental_deterministic = False ds = dataset_ops.Dataset.range(0).with_options(options1).with_options( options2) self.assertTrue(ds.options().experimental_autotune) # Explicitly check that flag is False since assertFalse allows None - self.assertIs(ds.options().experimental_filter_fusion, False) + self.assertIs(ds.options().experimental_deterministic, False) def testOptionsTwiceDifferentError(self): options1 = dataset_ops.Options() @@ -247,12 +247,12 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): options1 = dataset_ops.Options() options1.experimental_autotune = True options2 = dataset_ops.Options() - options2.experimental_filter_fusion = True + options2.experimental_deterministic = True ds = dataset_ops.Dataset.zip( (dataset_ops.Dataset.range(0).with_options(options1), dataset_ops.Dataset.range(0).with_options(options2))) self.assertTrue(ds.options().experimental_autotune) - self.assertTrue(ds.options().experimental_filter_fusion) + self.assertTrue(ds.options().experimental_deterministic) # TODO(b/119882922): use-after-free bug in eager mode. # pylint: disable=g-long-lambda diff --git a/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py b/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py index 886c9acc03be4485ba2e9f252d98faf61e8c16d3..622ebb55dec635c9b28787820ad789abb51a6fdf 100644 --- a/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py +++ b/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from tensorflow.core.protobuf import config_pb2 from tensorflow.python.data.experimental.ops import optimization +from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import multi_device_iterator_ops @@ -265,7 +266,8 @@ class MultiDeviceIteratorTest(test_base.DatasetTestBase): dataset = dataset.cache() options = dataset_ops.Options() - options.experimental_noop_elimination = True + options.experimental_optimization = OptimizationOptions() + options.experimental_optimization.noop_elimination = True dataset = dataset.with_options(options) multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator( diff --git a/tensorflow/python/data/ops/BUILD b/tensorflow/python/data/ops/BUILD index 27c9175ccbf8c6f6cef506c8fafaf7cd1ed57f71..0c5acda180f960d428266fc63df4cbafbdff6d7b 100644 --- a/tensorflow/python/data/ops/BUILD +++ b/tensorflow/python/data/ops/BUILD @@ -27,6 +27,7 @@ py_library( "//tensorflow/python:tensor_util", "//tensorflow/python:util", "//tensorflow/python/data/experimental/ops:filter_for_shard_ops", + "//tensorflow/python/data/experimental/ops:optimization_options", "//tensorflow/python/data/experimental/ops:stats_options", "//tensorflow/python/data/experimental/ops:threading_options", "//tensorflow/python/data/util:nest", diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index 51123aaf4430846c716642b0227aef5442eb7012..c57ddab15f2785ac03db5ae97ebdb7bfef2cbf73 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -27,6 +27,7 @@ import six from tensorflow.python.compat import compat from tensorflow.python.data.experimental.ops import filter_for_shard_ops +from tensorflow.python.data.experimental.ops import optimization_options from tensorflow.python.data.experimental.ops import stats_options from tensorflow.python.data.experimental.ops import threading_options from tensorflow.python.data.ops import iterator_ops @@ -1587,56 +1588,15 @@ class Options(options_lib.OptionsBase): "Whether to dynamically adjust the values of tunable parameters (e.g. " "degrees of parallelism).") - experimental_filter_fusion = options_lib.create_option( - name="experimental_filter_fusion", - ty=bool, - docstring="Whether to fuse filter transformations.") - - experimental_hoist_random_uniform = options_lib.create_option( - name="experimental_hoist_random_uniform", - ty=bool, - docstring= - "Whether to hoist `tf.random_uniform()` ops out of map transformations.") - - experimental_map_and_batch_fusion = options_lib.create_option( - name="experimental_map_and_batch_fusion", - ty=bool, - docstring="Whether to fuse map and batch transformations.") - - experimental_map_and_filter_fusion = options_lib.create_option( - name="experimental_map_and_filter_fusion", - ty=bool, - docstring="Whether to fuse map and filter transformations.") - - experimental_map_fusion = options_lib.create_option( - name="experimental_map_and_filter_fusion", - ty=bool, - docstring="Whether to fuse map transformations.") - - experimental_map_parallelization = options_lib.create_option( - name="experimental_map_parallelization", - ty=bool, - docstring="Whether to parallelize stateless map transformations.") - - experimental_map_vectorization = options_lib.create_option( - name="experimental_map_vectorization", - ty=bool, - docstring="Whether to vectorize map transformations.") - - experimental_noop_elimination = options_lib.create_option( - name="experimental_noop_elimination", - ty=bool, - docstring="Whether to eliminate no-op transformations.") - experimental_numa_aware = options_lib.create_option( name="experimental_numa_aware", ty=bool, docstring="Whether to use NUMA-aware operations.") - experimental_shuffle_and_repeat_fusion = options_lib.create_option( - name="experimental_shuffle_and_repeat_fusion", - ty=bool, - docstring="Whether to fuse shuffle and repeat transformations.") + experimental_optimization = options_lib.create_option( + name="experimental_optimization", + ty=optimization_options.OptimizationOptions, + docstring="Associates the given optimization options with the dataset.") experimental_stats = options_lib.create_option( name="experimental_stats", @@ -1650,29 +1610,30 @@ class Options(options_lib.OptionsBase): def _static_optimizations(self): """Produces the list of enabled static optimizations.""" - experimental_optimizations = [ - "filter_fusion", - "hoist_random_uniform", - "map_and_batch_fusion", - "map_and_filter_fusion", - "map_fusion", - "map_parallelization", - "map_vectorization", - "noop_elimination", - "shuffle_and_repeat_fusion", - ] - result = [] - for exp_opt in experimental_optimizations: - if getattr(self, "experimental_" + exp_opt): - result.append(exp_opt) - if getattr(self, "experimental_numa_aware"): + result = [] + exp_optimization_options = self.experimental_optimization + if exp_optimization_options: + optimizations = [ + "filter_fusion", + "hoist_random_uniform", + "map_and_batch_fusion", + "map_and_filter_fusion", + "map_fusion", + "map_parallelization", + "map_vectorization", + "noop_elimination", + "shuffle_and_repeat_fusion", + ] + for optimization in optimizations: + if getattr(exp_optimization_options, optimization): + result.append(optimization) + if self.experimental_numa_aware: result.append("make_numa_aware") - if getattr(self, "experimental_deterministic") is False: + if self.experimental_deterministic is False: result.append("make_sloppy") - experimental_stats_options = getattr(self, "experimental_stats") - if experimental_stats_options and getattr(experimental_stats_options, - "latency_all_edges"): + exp_stats_options = self.experimental_stats + if exp_stats_options and exp_stats_options.latency_all_edges: result.append("latency_all_edges") return result diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt index 024b4514ba7d6ad12e0de7d5f2cf7d1941b4263c..72fc2c3a9ee5b985723ce2dba9643ba796362dc7 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.data.-options.pbtxt @@ -11,44 +11,12 @@ tf_class { name: "experimental_deterministic" mtype: "" } - member { - name: "experimental_filter_fusion" - mtype: "" - } - member { - name: "experimental_hoist_random_uniform" - mtype: "" - } - member { - name: "experimental_map_and_batch_fusion" - mtype: "" - } - member { - name: "experimental_map_and_filter_fusion" - mtype: "" - } - member { - name: "experimental_map_fusion" - mtype: "" - } - member { - name: "experimental_map_parallelization" - mtype: "" - } - member { - name: "experimental_map_vectorization" - mtype: "" - } - member { - name: "experimental_noop_elimination" - mtype: "" - } member { name: "experimental_numa_aware" mtype: "" } member { - name: "experimental_shuffle_and_repeat_fusion" + name: "experimental_optimization" mtype: "" } member { diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..9ca75828e55cdaeac5a493f49fe4bd963265e9d4 --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt @@ -0,0 +1,46 @@ +path: "tensorflow.data.experimental.OptimizationOptions" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "filter_fusion" + mtype: "" + } + member { + name: "hoist_random_uniform" + mtype: "" + } + member { + name: "map_and_batch_fusion" + mtype: "" + } + member { + name: "map_and_filter_fusion" + mtype: "" + } + member { + name: "map_fusion" + mtype: "" + } + member { + name: "map_parallelization" + mtype: "" + } + member { + name: "map_vectorization" + mtype: "" + } + member { + name: "noop_elimination" + mtype: "" + } + member { + name: "shuffle_and_repeat_fusion" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt index a3cb799fc3d2a87877f77c15556feadf65e4baf8..ad10b8228334906aed1b71ad3387c7844e2a8e52 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt @@ -12,6 +12,10 @@ tf_module { name: "CsvDataset" mtype: "" } + member { + name: "OptimizationOptions" + mtype: "" + } member { name: "Optional" mtype: "" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt index 024b4514ba7d6ad12e0de7d5f2cf7d1941b4263c..72fc2c3a9ee5b985723ce2dba9643ba796362dc7 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.data.-options.pbtxt @@ -11,44 +11,12 @@ tf_class { name: "experimental_deterministic" mtype: "" } - member { - name: "experimental_filter_fusion" - mtype: "" - } - member { - name: "experimental_hoist_random_uniform" - mtype: "" - } - member { - name: "experimental_map_and_batch_fusion" - mtype: "" - } - member { - name: "experimental_map_and_filter_fusion" - mtype: "" - } - member { - name: "experimental_map_fusion" - mtype: "" - } - member { - name: "experimental_map_parallelization" - mtype: "" - } - member { - name: "experimental_map_vectorization" - mtype: "" - } - member { - name: "experimental_noop_elimination" - mtype: "" - } member { name: "experimental_numa_aware" mtype: "" } member { - name: "experimental_shuffle_and_repeat_fusion" + name: "experimental_optimization" mtype: "" } member { diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..9ca75828e55cdaeac5a493f49fe4bd963265e9d4 --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt @@ -0,0 +1,46 @@ +path: "tensorflow.data.experimental.OptimizationOptions" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "filter_fusion" + mtype: "" + } + member { + name: "hoist_random_uniform" + mtype: "" + } + member { + name: "map_and_batch_fusion" + mtype: "" + } + member { + name: "map_and_filter_fusion" + mtype: "" + } + member { + name: "map_fusion" + mtype: "" + } + member { + name: "map_parallelization" + mtype: "" + } + member { + name: "map_vectorization" + mtype: "" + } + member { + name: "noop_elimination" + mtype: "" + } + member { + name: "shuffle_and_repeat_fusion" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt index a3cb799fc3d2a87877f77c15556feadf65e4baf8..ad10b8228334906aed1b71ad3387c7844e2a8e52 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt @@ -12,6 +12,10 @@ tf_module { name: "CsvDataset" mtype: "" } + member { + name: "OptimizationOptions" + mtype: "" + } member { name: "Optional" mtype: ""