From 9b82752179bd4a61a9d33a638b8d9cb06adcf9e0 Mon Sep 17 00:00:00 2001 From: Derek Murray Date: Fri, 11 Oct 2019 08:09:07 -0700 Subject: [PATCH] Replace tf.{data.experimental,contrib.data}.unbatch() with Dataset.unbatch(). PiperOrigin-RevId: 274174768 --- .../data/experimental/benchmarks/unbatch_benchmark.py | 7 +++---- .../optimization/choose_fastest_branch_dataset_test.py | 5 ++--- .../choose_fastest_branch_dataset_serialization_test.py | 3 +-- .../serialization/unbatch_dataset_serialization_test.py | 3 +-- tensorflow/python/data/experimental/ops/batching.py | 2 +- tensorflow/python/data/experimental/ops/resampling.py | 3 +-- tensorflow/python/data/kernel_tests/unbatch_test.py | 2 +- tensorflow/python/keras/engine/training_distributed.py | 3 +-- tensorflow/python/tpu/datasets.py | 3 +-- 9 files changed, 12 insertions(+), 19 deletions(-) diff --git a/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py b/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py index 44ca6d9f39b..f27d812f0e0 100644 --- a/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py +++ b/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Benchmarks for `tf.data.experimental.unbatch()`.""" +"""Benchmarks for `tf.data.Dataset.unbatch()`.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function @@ -22,7 +22,6 @@ import time import numpy as np from tensorflow.python.client import session -from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -31,7 +30,7 @@ from tensorflow.python.platform import test class UnbatchBenchmark(test.Benchmark): - """Benchmarks for `tf.data.experimental.unbatch()`.""" + """Benchmarks for `tf.data.Dataset.unbatch()`.""" def benchmark_native_unbatch(self): batch_sizes = [1, 2, 5, 10, 20, 50] @@ -40,7 +39,7 @@ class UnbatchBenchmark(test.Benchmark): dataset = dataset_ops.Dataset.from_tensors("element").repeat(None) batch_size_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) dataset = dataset.batch(batch_size_placeholder) - dataset = dataset.apply(batching.unbatch()) + dataset = dataset.unbatch() dataset = dataset.skip(elems_per_trial) options = dataset_ops.Options() options.experimental_optimization.apply_default_optimizations = False diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/choose_fastest_branch_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/choose_fastest_branch_dataset_test.py index 10f3f340b50..ee05ae0603d 100644 --- a/tensorflow/python/data/experimental/kernel_tests/optimization/choose_fastest_branch_dataset_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/optimization/choose_fastest_branch_dataset_test.py @@ -19,7 +19,6 @@ from __future__ import print_function from absl.testing import parameterized -from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.experimental.ops import optimization from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops @@ -120,7 +119,7 @@ class ChooseFastestBranchDatasetTest(test_base.DatasetTestBase, dataset = dataset_ops.Dataset.from_tensors(0).repeat(1000).batch(100) def branch(dataset): - return dataset.apply(batching.unbatch()) + return dataset.unbatch() choose_fastest = optimization._ChooseFastestBranchDataset( dataset, [branch, branch], @@ -134,7 +133,7 @@ class ChooseFastestBranchDatasetTest(test_base.DatasetTestBase, dataset = dataset_ops.Dataset.from_tensors(0).repeat(1000).batch(100) def branch(dataset): - return dataset.apply(batching.unbatch()) + return dataset.unbatch() def make_dataset(): return optimization._ChooseFastestBranchDataset( diff --git a/tensorflow/python/data/experimental/kernel_tests/serialization/choose_fastest_branch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/choose_fastest_branch_dataset_serialization_test.py index 573b773a2e8..325dbc9187d 100644 --- a/tensorflow/python/data/experimental/kernel_tests/serialization/choose_fastest_branch_dataset_serialization_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/serialization/choose_fastest_branch_dataset_serialization_test.py @@ -20,7 +20,6 @@ from __future__ import print_function from absl.testing import parameterized from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base -from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.experimental.ops import optimization from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops @@ -99,7 +98,7 @@ class ChooseFastestBranchDatasetSerializationTest( dataset = dataset_ops.Dataset.from_tensors(0).repeat(1000).batch(100) def branch(dataset): - return dataset.apply(batching.unbatch()) + return dataset.unbatch() return optimization._ChooseFastestBranchDataset( dataset, [branch, branch], diff --git a/tensorflow/python/data/experimental/kernel_tests/serialization/unbatch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/unbatch_dataset_serialization_test.py index d788d4a9f03..6fa115bd2fe 100644 --- a/tensorflow/python/data/experimental/kernel_tests/serialization/unbatch_dataset_serialization_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/serialization/unbatch_dataset_serialization_test.py @@ -21,7 +21,6 @@ from absl.testing import parameterized import numpy as np from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base -from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import combinations @@ -39,7 +38,7 @@ class UnbatchDatasetSerializationTest( np.array(multiplier) * np.arange(tensor_slice_len)) return dataset_ops.Dataset.from_tensor_slices(components).batch( - batch_size).apply(batching.unbatch()) + batch_size).unbatch() @combinations.generate(test_base.default_test_combinations()) def testCore(self): diff --git a/tensorflow/python/data/experimental/ops/batching.py b/tensorflow/python/data/experimental/ops/batching.py index a550d77784e..e5dd2eee10c 100644 --- a/tensorflow/python/data/experimental/ops/batching.py +++ b/tensorflow/python/data/experimental/ops/batching.py @@ -270,7 +270,7 @@ def unbatch(): # of a dataset. a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] } - a.apply(tf.data.experimental.unbatch()) == { + a.unbatch() == { 'a', 'b', 'c', 'a', 'b', 'a', 'b', 'c', 'd'} ``` diff --git a/tensorflow/python/data/experimental/ops/resampling.py b/tensorflow/python/data/experimental/ops/resampling.py index 6676085ae59..a9da1a7d092 100644 --- a/tensorflow/python/data/experimental/ops/resampling.py +++ b/tensorflow/python/data/experimental/ops/resampling.py @@ -19,7 +19,6 @@ from __future__ import print_function import numpy as np -from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.experimental.ops import interleave_ops from tensorflow.python.data.experimental.ops import scan_ops from tensorflow.python.data.ops import dataset_ops @@ -182,7 +181,7 @@ def _estimate_initial_dist_ds( initial_dist_ds = (class_values_ds.batch(dist_estimation_batch_size) .apply(scan_ops.scan(initial_examples_per_class_seen, update_estimate_and_tile)) - .apply(batching.unbatch())) + .unbatch()) return initial_dist_ds diff --git a/tensorflow/python/data/kernel_tests/unbatch_test.py b/tensorflow/python/data/kernel_tests/unbatch_test.py index 8b766c44f7b..5bb4852d534 100644 --- a/tensorflow/python/data/kernel_tests/unbatch_test.py +++ b/tensorflow/python/data/kernel_tests/unbatch_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for `tf.data.experimental.unbatch()`.""" +"""Tests for `tf.data.Dataset.unbatch()`.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function diff --git a/tensorflow/python/keras/engine/training_distributed.py b/tensorflow/python/keras/engine/training_distributed.py index bce89663f7a..342ad133cd4 100644 --- a/tensorflow/python/keras/engine/training_distributed.py +++ b/tensorflow/python/keras/engine/training_distributed.py @@ -21,7 +21,6 @@ from __future__ import print_function import numpy as np -from tensorflow.python.data.experimental.ops import batching from tensorflow.python.distribute import distribute_coordinator as dc from tensorflow.python.distribute import distribution_strategy_context from tensorflow.python.distribute import input_lib @@ -457,7 +456,7 @@ def experimental_tpu_predict_loop(model, padding_handler.update_mask) dataset = dataset.map(padding_handler.pad_batch) - dataset = dataset.apply(batching.unbatch()) + dataset = dataset.unbatch() # Upon this point, it is guaranteed that the dataset does not # have partial batches. Thus, we set `drop_remainder=True` to # get static shape information about the elements in the dataset. diff --git a/tensorflow/python/tpu/datasets.py b/tensorflow/python/tpu/datasets.py index eab0249cc07..5f4c65f60dd 100644 --- a/tensorflow/python/tpu/datasets.py +++ b/tensorflow/python/tpu/datasets.py @@ -18,7 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.experimental.ops import interleave_ops from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops @@ -194,6 +193,6 @@ def StreamingFilesDataset(files, if batch_transfer_size: # Undo the batching used during the transfer. - output_dataset = output_dataset.apply(batching.unbatch()).prefetch(1) + output_dataset = output_dataset.unbatch().prefetch(1) return output_dataset