diff options
| author |  Jiri Simsa <jsimsa@google.com> | 2018-10-01 17:26:37 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-01 17:31:54 -0700 |
| commit | 9a23e9251ecba026471ff77a5bbbc802a2889a10 (patch) | |
| tree | b3cdc0a120ac823a91fbb8a8b68e0846f239885e /tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py | |
| parent | ea5c529ed7b7e17d1e66bc7cf4479d232ed0a896 (diff) | |
[tf.data] Adding `tf.data.Options()`, `tf.data.Dataset.options()`, and `tf.data.Dataset.with_options()` to make it possible to respectively represent, get, and set options, such as optimization configuration, of a tf.data input pipeline.
PiperOrigin-RevId: 215310764
Diffstat (limited to 'tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py')
| -rw-r--r-- | tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py | 14 |
1 files changed, 9 insertions, 5 deletions
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 d47492753e..32ebc49c40 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 @@ -69,10 +69,11 @@ class MapVectorizationTest(test_base.DatasetTestBase, parameterized.TestCase): map_fn, num_parallel_calls=num_parallel_calls).batch(batch_size) unoptimized = _make_dataset([map_node_name, "Batch"]) - optimized = _make_dataset(["Batch", map_node_name] if expect_optimized else - [map_node_name, "Batch"]).apply( - optimization.optimize(["map_vectorization"])) - + optimized = _make_dataset(["Batch", map_node_name] + if expect_optimized else [map_node_name, "Batch"]) + options = dataset_ops.Options() + options.experimental_map_vectorization = True + optimized = optimized.with_options(options) return unoptimized, optimized @parameterized.named_parameters( @@ -179,7 +180,10 @@ class MapVectorizationBenchmark(test.Benchmark): unoptimized = input_dataset.map(map_fn).batch(batch_size) unoptimized_op = unoptimized.make_one_shot_iterator().get_next() - optimized = unoptimized.apply(optimization.optimize(["map_vectorization"])) + optimized = input_dataset.map(map_fn).batch(batch_size) + options = dataset_ops.Options() + options.experimental_map_vectorization = True + optimized = optimized.with_options(options) optimized_op = optimized.make_one_shot_iterator().get_next() unoptimized_time = self._run( |