wehub-resource-sync
124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
"""Tests for utilities for traversing the dataset construction graph."""
|
|
|
|
from absl.testing import parameterized
|
|
|
|
from tensorflow.python.compat import compat
|
|
from tensorflow.python.data.experimental.ops import data_service_ops
|
|
from tensorflow.python.data.kernel_tests import test_base
|
|
from tensorflow.python.data.ops import dataset_ops
|
|
from tensorflow.python.data.util import traverse
|
|
from tensorflow.python.framework import combinations
|
|
from tensorflow.python.ops import gen_dataset_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.platform import test
|
|
|
|
|
|
class _TestDataset(dataset_ops.UnaryUnchangedStructureDataset):
|
|
|
|
def __init__(self, input_dataset):
|
|
self._input_dataset = input_dataset
|
|
temp_variant_tensor = gen_dataset_ops.prefetch_dataset(
|
|
input_dataset._variant_tensor,
|
|
buffer_size=1,
|
|
**self._flat_structure)
|
|
variant_tensor = gen_dataset_ops.model_dataset(
|
|
temp_variant_tensor, **self._flat_structure)
|
|
super(_TestDataset, self).__init__(input_dataset, variant_tensor)
|
|
|
|
|
|
class TraverseTest(test_base.DatasetTestBase, parameterized.TestCase):
|
|
|
|
@combinations.generate(test_base.graph_only_combinations())
|
|
def testOnlySource(self):
|
|
ds = dataset_ops.Dataset.range(10)
|
|
variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
|
|
self.assertAllEqual(["RangeDataset"], [x.name for x in variant_tensor_ops])
|
|
|
|
@combinations.generate(test_base.graph_only_combinations())
|
|
def testSimplePipeline(self):
|
|
ds = dataset_ops.Dataset.range(10).map(math_ops.square)
|
|
variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
|
|
self.assertSetEqual(
|
|
set(["MapDataset", "RangeDataset"]),
|
|
set(x.name for x in variant_tensor_ops))
|
|
|
|
@combinations.generate(test_base.graph_only_combinations())
|
|
def testConcat(self):
|
|
ds1 = dataset_ops.Dataset.range(10)
|
|
ds2 = dataset_ops.Dataset.range(10)
|
|
ds = ds1.concatenate(ds2)
|
|
variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
|
|
self.assertSetEqual(
|
|
set(["ConcatenateDataset", "RangeDataset", "RangeDataset_1"]),
|
|
set(x.name for x in variant_tensor_ops))
|
|
|
|
@combinations.generate(test_base.graph_only_combinations())
|
|
def testZip(self):
|
|
ds1 = dataset_ops.Dataset.range(10)
|
|
ds2 = dataset_ops.Dataset.range(10)
|
|
ds = dataset_ops.Dataset.zip((ds1, ds2))
|
|
variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
|
|
self.assertSetEqual(
|
|
set(["ZipDataset", "RangeDataset", "RangeDataset_1"]),
|
|
set(x.name for x in variant_tensor_ops))
|
|
|
|
@combinations.generate(test_base.graph_only_combinations())
|
|
def testMultipleVariantTensors(self):
|
|
ds = dataset_ops.Dataset.range(10)
|
|
ds = _TestDataset(ds)
|
|
variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
|
|
self.assertSetEqual(
|
|
set(["RangeDataset", "ModelDataset", "PrefetchDataset"]),
|
|
set(x.name for x in variant_tensor_ops))
|
|
|
|
@combinations.generate(test_base.graph_only_combinations())
|
|
def testFlatMap(self):
|
|
ds1 = dataset_ops.Dataset.range(10).repeat(10)
|
|
|
|
def map_fn(ds):
|
|
|
|
def _map(x):
|
|
return ds.batch(x)
|
|
|
|
return _map
|
|
|
|
ds2 = dataset_ops.Dataset.range(20).prefetch(1)
|
|
ds2 = ds2.flat_map(map_fn(ds1))
|
|
variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds2)
|
|
self.assertSetEqual(
|
|
set([
|
|
"FlatMapDataset", "PrefetchDataset", "RepeatDataset",
|
|
"RangeDataset", "RangeDataset_1"
|
|
]), set(x.name for x in variant_tensor_ops))
|
|
|
|
@combinations.generate(test_base.graph_only_combinations())
|
|
def testTfDataService(self):
|
|
ds = dataset_ops.Dataset.range(10)
|
|
ds = ds.apply(
|
|
data_service_ops.distribute("parallel_epochs", "grpc://foo:0"))
|
|
ops = traverse.obtain_capture_by_value_ops(ds)
|
|
data_service_dataset_op = ("DataServiceDatasetV4"
|
|
if compat.forward_compatible(2022, 8, 31) else
|
|
"DataServiceDatasetV3")
|
|
self.assertContainsSubset(
|
|
["RangeDataset", data_service_dataset_op, "DummyIterationCounter"],
|
|
set(x.name for x in ops))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test.main()
|