wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/data/experimental/ops/BUILD
+635
View File
@@ -0,0 +1,635 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
py_library(
name = "batching",
srcs = ["batching.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:structured_function",
"//tensorflow/python/data/util:convert",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "cardinality",
srcs = ["cardinality.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:dataset_ops_gen",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "compression_ops",
srcs = ["compression_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/util:structure",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
],
)
py_library(
name = "counter",
srcs = ["counter.py"],
strict_deps = True,
deps = [
"//tensorflow/python:tf2",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "data_service_ops",
srcs = [
"data_service_ops.py",
],
lazy_imports = True,
strict_deps = True,
deps = [
":compression_ops",
"//tensorflow/core/protobuf:for_core_protos_py_proto",
"//tensorflow/python:tf2",
"//tensorflow/python/data/experimental/service:_pywrap_server_lib",
"//tensorflow/python/data/experimental/service:_pywrap_utils_exp",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/data/ops:structured_function",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/saved_model:nested_structure_coder",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "data_service_ops_test",
srcs = ["data_service_ops_test.py"],
deps = [
":data_service_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "distributed_save_op",
srcs = [
"distributed_save_op.py",
],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/saved_model:nested_structure_coder",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "distribute",
srcs = [
"distribute.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/types:data",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "enumerate_ops",
srcs = ["enumerate_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "error_ops",
srcs = ["error_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "from_list",
srcs = ["from_list.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/data/util:structure",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "get_single_element",
srcs = ["get_single_element.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/types:data",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "global_shuffle_op",
srcs = [
"global_shuffle_op.py",
],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/data/util:random_seed",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:dataset_ops_gen",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
],
)
py_library(
name = "grouping",
srcs = ["grouping.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:structured_function",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/data/util:structure",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "index_flat_map_op",
srcs = [
"index_flat_map_op.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:structured_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
],
)
py_library(
name = "interleave_ops",
srcs = ["interleave_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python:tf2",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:readers",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "io",
srcs = [
"io.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "iterator_ops",
srcs = [
"iterator_ops.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:proto_ops",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "iterator_model_ops",
srcs = [
"iterator_model_ops.py",
],
strict_deps = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "lookup_ops",
srcs = [
"lookup_ops.py",
],
strict_deps = True,
deps = [
":cardinality",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/ops:lookup_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "map_defun",
srcs = ["map_defun.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:dataset_ops_gen",
],
)
py_library(
name = "matching_files",
srcs = ["matching_files.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
],
)
py_library(
name = "pad_to_cardinality",
srcs = ["pad_to_cardinality.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/eager:context",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "parsing_ops",
srcs = ["parsing_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/util:structure",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/ops:parsing_ops",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "prefetching_ops",
srcs = ["prefetching_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/data/ops:structured_function",
"//tensorflow/python/data/util:structure",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:function",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:dataset_ops_gen",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/ops:functional_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "random_access",
srcs = ["random_access.py"],
strict_deps = True,
deps = [
"//tensorflow/python/data/util:structure",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "random_ops",
srcs = [
"random_ops.py",
],
strict_deps = True,
deps = [
"//tensorflow/python:tf2",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "readers",
srcs = [
"readers.py",
],
strict_deps = True,
deps = [
":error_ops",
":parsing_ops",
"//tensorflow/python:tf2",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/data/ops:readers",
"//tensorflow/python/data/util:convert",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/ops:io_ops",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "resampling",
srcs = ["resampling.py"],
strict_deps = True,
deps = [
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "scan_ops",
srcs = ["scan_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "shuffle_ops",
srcs = [
"shuffle_ops.py",
],
strict_deps = True,
visibility = [
# TODO(jsimsa): Remove non-default visibility exception when `index_shuffle` is made public.
"//third_party/py/tensorflow_datasets:__subpackages__",
"//tensorflow:internal",
],
deps = [
":random_access",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/util:random_seed",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:dataset_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "snapshot",
srcs = [
"snapshot.py",
],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "take_while_ops",
srcs = ["take_while_ops.py"],
strict_deps = True,
deps = [
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "testing",
srcs = ["testing.py"],
strict_deps = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
],
)
py_library(
name = "unique",
srcs = [
"unique.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "weighted_flat_map_op",
srcs = [
"weighted_flat_map_op.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/util:convert",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/data/util:random_seed",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:dataset_ops_gen",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
],
)
py_library(
name = "writers",
srcs = [
"writers.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/util:convert",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:experimental_dataset_ops_gen",
"//tensorflow/python/types:data",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "dataset_ops",
strict_deps = True,
deps = [
":batching",
":cardinality",
":compression_ops",
":counter",
":data_service_ops",
":distribute",
":distributed_save_op",
":enumerate_ops",
":error_ops",
":from_list",
":get_single_element",
":global_shuffle_op",
":grouping",
":interleave_ops",
":io",
":map_defun",
":matching_files",
":pad_to_cardinality",
":prefetching_ops",
":random_access",
":readers",
":resampling",
":scan_ops",
":shuffle_ops",
":snapshot",
":take_while_ops",
":unique",
":writers",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:debug_mode",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/ops:dataset_ops_gen",
],
)
tensorflow/python/data/experimental/ops/batching.py
+379
View File
@@ -0,0 +1,379 @@
# Copyright 2017 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.
# ==============================================================================
"""Batching dataset transformations."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import structured_function
from tensorflow.python.data.util import convert
from tensorflow.python.data.util import nest
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.dense_to_ragged_batch")
@deprecation.deprecated(None, "Use `tf.data.Dataset.ragged_batch` instead.")
def dense_to_ragged_batch(batch_size,
drop_remainder=False,
row_splits_dtype=dtypes.int64):
"""A transformation that batches ragged elements into `tf.RaggedTensor`s.
This transformation combines multiple consecutive elements of the input
dataset into a single element.
Like `tf.data.Dataset.batch`, the components of the resulting element will
have an additional outer dimension, which will be `batch_size` (or
`N % batch_size` for the last element if `batch_size` does not divide the
number of input elements `N` evenly and `drop_remainder` is `False`). If
your program depends on the batches having the same outer dimension, you
should set the `drop_remainder` argument to `True` to prevent the smaller
batch from being produced.
Unlike `tf.data.Dataset.batch`, the input elements to be batched may have
different shapes:
* If an input element is a `tf.Tensor` whose static `tf.TensorShape` is
fully defined, then it is batched as normal.
* If an input element is a `tf.Tensor` whose static `tf.TensorShape` contains
one or more axes with unknown size (i.e., `shape[i]=None`), then the output
will contain a `tf.RaggedTensor` that is ragged up to any of such
dimensions.
* If an input element is a `tf.RaggedTensor` or any other type, then it is
batched as normal.
Example:
>>> dataset = tf.data.Dataset.from_tensor_slices(np.arange(6))
>>> dataset = dataset.map(lambda x: tf.range(x))
>>> dataset.element_spec.shape
TensorShape([None])
>>> dataset = dataset.apply(
... tf.data.experimental.dense_to_ragged_batch(batch_size=2))
>>> for batch in dataset:
... print(batch)
<tf.RaggedTensor [[], [0]]>
<tf.RaggedTensor [[0, 1], [0, 1, 2]]>
<tf.RaggedTensor [[0, 1, 2, 3], [0, 1, 2, 3, 4]]>
Args:
batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
consecutive elements of this dataset to combine in a single batch.
drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
whether the last batch should be dropped in the case it has fewer than
`batch_size` elements; the default behavior is not to drop the smaller
batch.
row_splits_dtype: The dtype that should be used for the `row_splits` of any
new ragged tensors. Existing `tf.RaggedTensor` elements do not have their
row_splits dtype changed.
Returns:
Dataset: A `Dataset`.
"""
def _apply_fn(dataset):
return dataset.ragged_batch(batch_size, drop_remainder, row_splits_dtype)
return _apply_fn
@tf_export("data.experimental.dense_to_sparse_batch")
@deprecation.deprecated(None, "Use `tf.data.Dataset.sparse_batch` instead.")
def dense_to_sparse_batch(batch_size, row_shape):
"""A transformation that batches ragged elements into `tf.sparse.SparseTensor`s.
Like `Dataset.padded_batch()`, this transformation combines multiple
consecutive elements of the dataset, which might have different
shapes, into a single element. The resulting element has three
components (`indices`, `values`, and `dense_shape`), which
comprise a `tf.sparse.SparseTensor` that represents the same data. The
`row_shape` represents the dense shape of each row in the
resulting `tf.sparse.SparseTensor`, to which the effective batch size is
prepended. For example:
```python
# NOTE: The following examples use `{ ... }` to represent the
# contents of a dataset.
a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] }
a.apply(tf.data.experimental.dense_to_sparse_batch(
batch_size=2, row_shape=[6])) ==
{
([[0, 0], [0, 1], [0, 2], [1, 0], [1, 1]], # indices
['a', 'b', 'c', 'a', 'b'], # values
[2, 6]), # dense_shape
([[0, 0], [0, 1], [0, 2], [0, 3]],
['a', 'b', 'c', 'd'],
[1, 6])
}
```
Args:
batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
consecutive elements of this dataset to combine in a single batch.
row_shape: A `tf.TensorShape` or `tf.int64` vector tensor-like object
representing the equivalent dense shape of a row in the resulting
`tf.sparse.SparseTensor`. Each element of this dataset must have the same
rank as `row_shape`, and must have size less than or equal to `row_shape`
in each dimension.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.sparse_batch(batch_size, row_shape)
return _apply_fn
@deprecation.deprecated(None, "Use `tf.data.experimental.map_and_batch()")
@tf_export(v1=["data.experimental.map_and_batch_with_legacy_function"])
def map_and_batch_with_legacy_function(map_func,
batch_size,
num_parallel_batches=None,
drop_remainder=False,
num_parallel_calls=None):
"""Fused implementation of `map` and `batch`.
NOTE: This is an escape hatch for existing uses of `map_and_batch` that do not
work with V2 functions. New uses are strongly discouraged and existing uses
should migrate to `map_and_batch` as this method will not be removed in V2.
Args:
map_func: A function mapping a nested structure of tensors to another
nested structure of tensors.
batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
consecutive elements of this dataset to combine in a single batch.
num_parallel_batches: (Optional.) A `tf.int64` scalar `tf.Tensor`,
representing the number of batches to create in parallel. On one hand,
higher values can help mitigate the effect of stragglers. On the other
hand, higher values can increase contention if CPU is scarce.
drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
whether the last batch should be dropped in case its size is smaller than
desired; the default behavior is not to drop the smaller batch.
num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
representing the number of elements to process in parallel. If not
specified, `batch_size * num_parallel_batches` elements will be processed
in parallel. If the value `tf.data.AUTOTUNE` is used, then
the number of parallel calls is set dynamically based on available CPU.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
Raises:
ValueError: If both `num_parallel_batches` and `num_parallel_calls` are
specified.
"""
if num_parallel_batches is None and num_parallel_calls is None:
num_parallel_calls = batch_size
elif num_parallel_batches is not None and num_parallel_calls is None:
num_parallel_calls = batch_size * num_parallel_batches
elif num_parallel_batches is not None and num_parallel_calls is not None:
raise ValueError(
"`map_and_batch_with_legacy_function` allows only one of "
"`num_parallel_batches` and "
"`num_parallel_calls` to be set, but "
f"`num_parallel_batches` was set to {num_parallel_batches} "
f"and `num_parallel_calls` as set to {num_parallel_calls}.")
def _apply_fn(dataset):
return _MapAndBatchDataset(dataset, map_func, batch_size,
num_parallel_calls, drop_remainder,
use_legacy_function=True)
return _apply_fn
@deprecation.deprecated(
None,
"Use `tf.data.Dataset.map(map_func, num_parallel_calls)` followed by "
"`tf.data.Dataset.batch(batch_size, drop_remainder)`. Static tf.data "
"optimizations will take care of using the fused implementation.")
@tf_export("data.experimental.map_and_batch")
def map_and_batch(map_func,
batch_size,
num_parallel_batches=None,
drop_remainder=False,
num_parallel_calls=None):
"""Fused implementation of `map` and `batch`.
Maps `map_func` across `batch_size` consecutive elements of this dataset
and then combines them into a batch. Functionally, it is equivalent to `map`
followed by `batch`. This API is temporary and deprecated since input pipeline
optimization now fuses consecutive `map` and `batch` operations automatically.
Args:
map_func: A function mapping a nested structure of tensors to another
nested structure of tensors.
batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
consecutive elements of this dataset to combine in a single batch.
num_parallel_batches: (Optional.) A `tf.int64` scalar `tf.Tensor`,
representing the number of batches to create in parallel. On one hand,
higher values can help mitigate the effect of stragglers. On the other
hand, higher values can increase contention if CPU is scarce.
drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
whether the last batch should be dropped in case its size is smaller than
desired; the default behavior is not to drop the smaller batch.
num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
representing the number of elements to process in parallel. If not
specified, `batch_size * num_parallel_batches` elements will be processed
in parallel. If the value `tf.data.AUTOTUNE` is used, then
the number of parallel calls is set dynamically based on available CPU.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
Raises:
ValueError: If both `num_parallel_batches` and `num_parallel_calls` are
specified.
"""
if num_parallel_batches is None and num_parallel_calls is None:
num_parallel_calls = batch_size
elif num_parallel_batches is not None and num_parallel_calls is None:
num_parallel_calls = batch_size * num_parallel_batches
elif num_parallel_batches is not None and num_parallel_calls is not None:
raise ValueError(
"`map_and_batch` allows only one of `num_parallel_batches` and "
"`num_parallel_calls` to be set, but "
f"`num_parallel_batches` was set to {num_parallel_batches} "
f"and `num_parallel_calls` as set to {num_parallel_calls}.")
def _apply_fn(dataset):
return _MapAndBatchDataset(dataset, map_func, batch_size,
num_parallel_calls, drop_remainder)
return _apply_fn
@deprecation.deprecated(None, "Use `tf.data.Dataset.unbatch()`.")
@tf_export("data.experimental.unbatch")
def unbatch():
"""Splits elements of a dataset into multiple elements on the batch dimension.
For example, if elements of the dataset are shaped `[B, a0, a1, ...]`,
where `B` may vary for each input element, then for each element in the
dataset, the unbatched dataset will contain `B` consecutive elements
of shape `[a0, a1, ...]`.
```python
# NOTE: The following example uses `{ ... }` to represent the contents
# of a dataset.
a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] }
a.unbatch() == {
'a', 'b', 'c', 'a', 'b', 'a', 'b', 'c', 'd'}
```
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.unbatch()
return _apply_fn
class _DenseToSparseBatchDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that batches ragged dense elements into `tf.sparse.SparseTensor`s."""
def __init__(self, input_dataset, batch_size, row_shape):
"""See `Dataset.dense_to_sparse_batch()` for more details."""
if not isinstance(
dataset_ops.get_legacy_output_types(input_dataset), dtypes.DType):
raise TypeError("`dense_to_sparse_batch` requires an input dataset whose "
"elements have a single component, but the given dataset "
"has the following component types: "
f"{dataset_ops.get_legacy_output_types(input_dataset)}.")
self._input_dataset = input_dataset
self._batch_size = batch_size
self._row_shape = row_shape
self._element_spec = sparse_tensor.SparseTensorSpec(
tensor_shape.TensorShape([None]).concatenate(self._row_shape),
dataset_ops.get_legacy_output_types(input_dataset))
variant_tensor = ged_ops.dense_to_sparse_batch_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._batch_size,
row_shape=convert.partial_shape_to_tensor(self._row_shape),
**self._flat_structure)
super(_DenseToSparseBatchDataset, self).__init__(input_dataset,
variant_tensor)
@property
def element_spec(self):
return self._element_spec
class _MapAndBatchDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that maps a function over a batch of elements."""
def __init__(self, input_dataset, map_func, batch_size, num_parallel_calls,
drop_remainder, use_legacy_function=False):
self._input_dataset = input_dataset
self._map_func = structured_function.StructuredFunctionWrapper(
map_func,
"tf.data.experimental.map_and_batch()",
dataset=input_dataset,
use_legacy_function=use_legacy_function)
self._batch_size_t = ops.convert_to_tensor(
batch_size, dtype=dtypes.int64, name="batch_size")
self._num_parallel_calls_t = ops.convert_to_tensor(
num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
self._drop_remainder_t = ops.convert_to_tensor(
drop_remainder, dtype=dtypes.bool, name="drop_remainder")
constant_drop_remainder = tensor_util.constant_value(self._drop_remainder_t)
# pylint: disable=protected-access
if constant_drop_remainder:
# NOTE(mrry): `constant_drop_remainder` may be `None` (unknown statically)
# or `False` (explicitly retaining the remainder).
# pylint: disable=g-long-lambda
self._element_spec = nest.map_structure(
lambda component_spec: component_spec._batch(
tensor_util.constant_value(self._batch_size_t)),
self._map_func.output_structure)
else:
self._element_spec = nest.map_structure(
lambda component_spec: component_spec._batch(None),
self._map_func.output_structure)
# pylint: enable=protected-access
variant_tensor = ged_ops.map_and_batch_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._map_func.function.captured_inputs,
f=self._map_func.function,
batch_size=self._batch_size_t,
num_parallel_calls=self._num_parallel_calls_t,
drop_remainder=self._drop_remainder_t,
preserve_cardinality=True,
**self._flat_structure)
super(_MapAndBatchDataset, self).__init__(input_dataset, variant_tensor)
def _functions(self):
return [self._map_func]
@property
def element_spec(self):
return self._element_spec
tensorflow/python/data/experimental/ops/cardinality.py
+113
View File
@@ -0,0 +1,113 @@
# 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.
# ==============================================================================
"""Cardinality analysis of `Dataset` objects."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_dataset_ops
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.util.tf_export import tf_export
INFINITE = -1
UNKNOWN = -2
tf_export("data.experimental.INFINITE_CARDINALITY").export_constant(
__name__, "INFINITE")
tf_export("data.experimental.UNKNOWN_CARDINALITY").export_constant(
__name__, "UNKNOWN")
# TODO(b/157691652): Deprecate this method after migrating users to the new API.
@tf_export("data.experimental.cardinality")
def cardinality(dataset):
"""Returns the cardinality of `dataset`, if known.
The operation returns the cardinality of `dataset`. The operation may return
`tf.data.experimental.INFINITE_CARDINALITY` if `dataset` contains an infinite
number of elements or `tf.data.experimental.UNKNOWN_CARDINALITY` if the
analysis fails to determine the number of elements in `dataset` (e.g. when the
dataset source is a file).
>>> dataset = tf.data.Dataset.range(42)
>>> print(tf.data.experimental.cardinality(dataset).numpy())
42
>>> dataset = dataset.repeat()
>>> cardinality = tf.data.experimental.cardinality(dataset)
>>> print((cardinality == tf.data.experimental.INFINITE_CARDINALITY).numpy())
True
>>> dataset = dataset.filter(lambda x: True)
>>> cardinality = tf.data.experimental.cardinality(dataset)
>>> print((cardinality == tf.data.experimental.UNKNOWN_CARDINALITY).numpy())
True
Args:
dataset: A `tf.data.Dataset` for which to determine cardinality.
Returns:
A scalar `tf.int64` `Tensor` representing the cardinality of `dataset`. If
the cardinality is infinite or unknown, the operation returns the named
constant `INFINITE_CARDINALITY` and `UNKNOWN_CARDINALITY` respectively.
"""
return gen_dataset_ops.dataset_cardinality(dataset._variant_tensor) # pylint: disable=protected-access
@tf_export("data.experimental.assert_cardinality")
def assert_cardinality(expected_cardinality):
"""Asserts the cardinality of the input dataset.
NOTE: The following assumes that "examples.tfrecord" contains 42 records.
>>> dataset = tf.data.TFRecordDataset("examples.tfrecord")
>>> cardinality = tf.data.experimental.cardinality(dataset)
>>> print((cardinality == tf.data.experimental.UNKNOWN_CARDINALITY).numpy())
True
>>> dataset = dataset.apply(tf.data.experimental.assert_cardinality(42))
>>> print(tf.data.experimental.cardinality(dataset).numpy())
42
Args:
expected_cardinality: The expected cardinality of the input dataset.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
Raises:
FailedPreconditionError: The assertion is checked at runtime (when iterating
the dataset) and an error is raised if the actual and expected cardinality
differ.
"""
def _apply_fn(dataset):
return _AssertCardinalityDataset(dataset, expected_cardinality)
return _apply_fn
class _AssertCardinalityDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that assert the cardinality of its input."""
def __init__(self, input_dataset, expected_cardinality):
self._input_dataset = input_dataset
self._expected_cardinality = ops.convert_to_tensor(
expected_cardinality, dtype=dtypes.int64, name="expected_cardinality")
# pylint: enable=protected-access
variant_tensor = ged_ops.assert_cardinality_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._expected_cardinality,
**self._flat_structure)
super(_AssertCardinalityDataset, self).__init__(input_dataset,
variant_tensor)
tensorflow/python/data/experimental/ops/compression_ops.py
+51
View File
@@ -0,0 +1,51 @@
# Copyright 2020 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.
# ==============================================================================
"""Ops for compressing and uncompressing dataset elements."""
from tensorflow.python.data.util import structure
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
def compress(element):
"""Compress a dataset element.
Args:
element: A nested structure of types supported by Tensorflow.
Returns:
A variant tensor representing the compressed element. This variant can be
passed to `uncompress` to get back the original element.
"""
element_spec = structure.type_spec_from_value(element)
tensor_list = structure.to_tensor_list(element_spec, element)
return ged_ops.compress_element(tensor_list)
def uncompress(element, output_spec):
"""Uncompress a compressed dataset element.
Args:
element: A scalar variant tensor to uncompress. The element should have been
created by calling `compress`.
output_spec: A nested structure of `tf.TypeSpec` representing the type(s) of
the uncompressed element.
Returns:
The uncompressed element.
"""
flat_types = structure.get_flat_tensor_types(output_spec)
flat_shapes = structure.get_flat_tensor_shapes(output_spec)
tensor_list = ged_ops.uncompress_element(
element, output_types=flat_types, output_shapes=flat_shapes)
return structure.from_tensor_list(output_spec, tensor_list)
tensorflow/python/data/experimental/ops/counter.py
+84
View File
@@ -0,0 +1,84 @@
# Copyright 2017 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.
# ==============================================================================
"""The Counter Dataset."""
from tensorflow.python import tf2
from tensorflow.python.compat import v2_compat
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.Counter", v1=[])
@deprecation.deprecated(None, "Use `tf.data.Dataset.counter(...)` instead.")
def CounterV2(start=0, step=1, dtype=dtypes.int64):
"""Creates a `Dataset` that counts from `start` in steps of size `step`.
Unlike `tf.data.Dataset.range` which will stop at some ending number,
`Counter` will produce elements indefinitely.
>>> dataset = tf.data.experimental.Counter().take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[0, 1, 2, 3, 4]
>>> dataset.element_spec
TensorSpec(shape=(), dtype=tf.int64, name=None)
>>> dataset = tf.data.experimental.Counter(dtype=tf.int32)
>>> dataset.element_spec
TensorSpec(shape=(), dtype=tf.int32, name=None)
>>> dataset = tf.data.experimental.Counter(start=2).take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[2, 3, 4, 5, 6]
>>> dataset = tf.data.experimental.Counter(start=2, step=5).take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[2, 7, 12, 17, 22]
>>> dataset = tf.data.experimental.Counter(start=10, step=-1).take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[10, 9, 8, 7, 6]
Args:
start: (Optional.) The starting value for the counter. Defaults to 0.
step: (Optional.) The step size for the counter. Defaults to 1.
dtype: (Optional.) The data type for counter elements. Defaults to
`tf.int64`.
Returns:
A `Dataset` of scalar `dtype` elements.
"""
return dataset_ops.Dataset.counter(start, step, dtype)
@tf_export(v1=["data.experimental.Counter"])
@deprecation.deprecated(None, "Use `tf.data.Dataset.counter(...)` instead.")
def CounterV1(start=0, step=1, dtype=dtypes.int64):
return dataset_ops.DatasetV1Adapter(CounterV2(start, step, dtype))
CounterV1.__doc__ = CounterV2.__doc__
if tf2.enabled():
Counter = CounterV2
else:
Counter = CounterV1
def _tf2_callback(): # pylint: disable=invalid-name
global Counter
if tf2.enabled():
Counter = CounterV2
else:
Counter = CounterV1
v2_compat.register_data_v2_callback(_tf2_callback)
tensorflow/python/data/experimental/ops/data_service_ops.py
+1231
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/data/experimental/ops/data_service_ops_test.py
+99
View File
@@ -0,0 +1,99 @@
# Copyright 2024 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 HashableElementSpec and related functionality."""
from tensorflow.python.data.experimental.ops import data_service_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import test
class HashableElementSpecTest(test.TestCase):
def testEqual(self):
spec1 = data_service_ops.HashableElementSpec(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
)
spec2 = data_service_ops.HashableElementSpec(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
)
self.assertEqual(spec1, spec2)
self.assertEqual(hash(spec1), hash(spec2))
def testNotEqual(self):
spec1 = data_service_ops.HashableElementSpec(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
)
spec2 = data_service_ops.HashableElementSpec(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.int64)
)
self.assertNotEqual(spec1, spec2)
self.assertNotEqual(hash(spec1), hash(spec2))
def testNotEqualOtherType(self):
spec1 = data_service_ops.HashableElementSpec(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
)
self.assertNotEqual(spec1, 123)
def testGetUncompressFuncCache(self):
spec1 = tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
spec2 = tensor_spec.TensorSpec(shape=(), dtype=dtypes.int64)
data_service_ops._get_uncompress_func.cache_clear()
func1 = data_service_ops._get_uncompress_func(
data_service_ops.HashableElementSpec(spec1)
)
func2 = data_service_ops._get_uncompress_func(
data_service_ops.HashableElementSpec(spec1)
)
func3 = data_service_ops._get_uncompress_func(
data_service_ops.HashableElementSpec(spec2)
)
self.assertIs(func1, func2)
self.assertIsNot(func1, func3)
def testRaggedTensorUncompressFuncCache(self):
spec1 = ragged_tensor.RaggedTensorSpec(
[3, None], dtypes.int32, 1, dtypes.int64
)
spec1_long_format = ragged_tensor.RaggedTensorSpec(
tensor_shape.TensorShape(
[tensor_shape.Dimension(3), tensor_shape.Dimension(None)]
),
dtypes.int32,
1,
dtypes.int64,
)
spec2 = ragged_tensor.RaggedTensorSpec(
([3, None]), dtypes.int64, 1, dtypes.int64
)
data_service_ops._get_uncompress_func.cache_clear()
func1a = data_service_ops._get_uncompress_func(
data_service_ops.HashableElementSpec(spec1)
)
func1b = data_service_ops._get_uncompress_func(
data_service_ops.HashableElementSpec(spec1_long_format)
)
func2 = data_service_ops._get_uncompress_func(
data_service_ops.HashableElementSpec(spec2)
)
self.assertIs(func1a, func1b)
self.assertIsNot(func1a, func2)
if __name__ == "__main__":
test.main()
tensorflow/python/data/experimental/ops/distribute.py
+399
View File
@@ -0,0 +1,399 @@
# Copyright 2019 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.
# ==============================================================================
"""Distribution Strategy-related dataset transformations."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops.options import ExternalStatePolicy
from tensorflow.python.data.util import nest
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.types import data as data_types
from tensorflow.python.util.tf_export import tf_export
SHARD_HINT = -1
tf_export("data.experimental.SHARD_HINT").export_constant(
__name__, "SHARD_HINT")
class _AutoShardDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that shards the `Dataset` automatically.
This dataset takes in an existing dataset and tries to automatically figure
out how to shard the dataset in a multi-worker scenario using graph rewrites.
If the AutoShardPolicy is set to FILE, it walks up the dataset graph until
it finds a reader dataset, then inserts a ShardDataset op before that node
so that each worker only sees some files.
If the AutoShardPolicy is set to DATA, it inserts a ShardDataset op at the
end of the input pipeline, before any terminal PrefetchDataset if there is
one. Additionally, if there is a RebatchDatasetV2 in the input pipeline, it
is written to legacy RebatchDataset for correctness reasons, since
RebatchDatasetV2 is incompatible with data sharding.
If the AutoShardPolicy is set to AUTO, it tries to do file-based sharding.
If it cannot find a reader dataset, it falls back to doing data-based
sharding.
If the AutoShardPolicy is set to OFF, it does nothing.
Attributes:
num_workers: Total number of workers to shard this dataset across.
index: The current worker index (out of the total number of workers) this
dataset is for.
num_replicas: The total number of replicas across all workers. This is used
only when sharding by data (either DATA or AUTO) in order to rewrite
RebatchDatasetV2 to RebatchDataset.
Raises:
NotFoundError: If we cannot find a suitable reader dataset to begin
automatically sharding the dataset.
"""
def __init__(self, input_dataset, num_workers, index, num_replicas=None):
self._input_dataset = input_dataset
self._element_spec = input_dataset.element_spec
variant_tensor = ged_ops.auto_shard_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
num_workers=num_workers,
index=index,
auto_shard_policy=int(
input_dataset.options().experimental_distribute.auto_shard_policy),
num_replicas=num_replicas,
**self._flat_structure)
super(_AutoShardDataset, self).__init__(input_dataset, variant_tensor)
@property
def element_spec(self):
return self._element_spec
def _AutoShardDatasetV1(input_dataset, num_workers, index, num_replicas=None): # pylint: disable=invalid-name
return dataset_ops.DatasetV1Adapter(
_AutoShardDataset(input_dataset, num_workers, index, num_replicas))
class _LegacyRebatchDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that divides its input batches into `num_replicas` sub-batches.
For each batch in the input dataset, _LegacyRebatchDataset will produce
`num_replicas` smaller batches whose sizes add up to the original batch size.
For example:
```python
ds = tf.data.Dataset.range(8)
ds = ds.batch(4)
ds = _LegacyRebatchDataset(ds, num_replicas=3)
for elem in ds:
print(elem)
>> [0, 1], [2, 3], [], [4, 5], [6, 7], []
```
"""
def __init__(self, input_dataset, num_replicas):
"""Creates a _LegacyRebatchDataset.
Args:
input_dataset: `Dataset` to rebatch.
num_replicas: A `tf.int64` scalar, representing the number of sub-batches
to split each batch from `input_dataset` into.
"""
def recalculate_batch_size(type_spec):
"""Recalculates the output_shape after dividing it by num_replicas."""
output_shape = type_spec._to_legacy_output_shapes() # pylint: disable=protected-access
if not isinstance(output_shape, tensor_shape.TensorShape):
return None
# If the output shape is unknown, we set the batch dimension to unknown.
if output_shape.rank is None:
return None
if len(output_shape) < 1:
raise ValueError(
"Invalid `input_dataset`. Expected a dataset whose elements "
"have rank >= 1 but found a dataset whose elements are scalars. "
"Fix the issue by adding the `batch` transformation to the "
"dataset.")
output_dims = [d.value for d in output_shape.dims]
if output_dims[0] is not None and output_dims[0] % num_replicas == 0:
return output_dims[0] // num_replicas
# Set the batch dimension to unknown. If the global batch size does not
# divide num_replicas evenly, the minibatches may have different sizes.
return None
def rebatch(type_spec):
# pylint: disable=protected-access
batch_size = recalculate_batch_size(type_spec)
return type_spec._unbatch()._batch(batch_size)
# pylint: enable=protected-access
self._element_spec = nest.map_structure(
rebatch, dataset_ops.get_structure(input_dataset))
# auto_shard rewrite assumes that there's normalize_to_dense before
# rebatch_dataset.
# LINT.IfChange
input_dataset = dataset_ops.normalize_to_dense(input_dataset)
variant_tensor = ged_ops.rebatch_dataset(
input_dataset._variant_tensor, # pylint: disable=protected-access
num_replicas=num_replicas,
**self._flat_structure)
# LINT.ThenChange(//tensorflow/core/grappler/optimizers/data/auto_shard.cc)
super(_LegacyRebatchDataset, self).__init__(input_dataset, variant_tensor)
@property
def element_spec(self):
return self._element_spec
class _RemoteDataset(dataset_ops.DatasetSource):
"""Creates a dataset on a given `device` given a graph def."""
def __init__(self, graph_def, device, element_spec):
self._elem_spec = element_spec
with ops.device(device):
variant_tensor = ged_ops.dataset_from_graph(graph_def)
super(_RemoteDataset, self).__init__(variant_tensor)
@property
def element_spec(self):
return self._elem_spec
def replicate(dataset, devices):
"""A transformation that replicates `dataset` onto a list of devices.
Args:
dataset: A `tf.data.Dataset` object.
devices: A list of devices to replicate the dataset on.
Returns:
A dictionary mapping device name to a dataset on that device.
"""
if not isinstance(dataset, data_types.DatasetV2):
raise TypeError(
f"Invalid `dataset`. Expected a `tf.data.Dataset` object but "
f"got {type(dataset)}.")
# pylint: disable=protected-access
dataset_device = dataset._variant_tensor.device
datasets = {}
if len(devices) == 1 and devices[0] == dataset_device:
datasets[devices[0]] = dataset
return datasets
with ops.colocate_with(dataset._variant_tensor):
dataset = dataset._apply_debug_options()
graph_def = dataset._as_serialized_graph(
strip_device_assignment=True,
external_state_policy=ExternalStatePolicy.WARN)
for device in devices:
ds = _RemoteDataset(graph_def, device, dataset.element_spec)
datasets[device] = ds
return datasets
def batch_sizes_for_worker(global_batch_size, num_workers,
num_replicas_per_worker, worker_index):
"""Determines how to rebatch a dataset for the given worker.
Given the global batch size, number of workers, number of replicas per worker,
and worker index, returns the correct batch sizes for rebatching a dataset
on worker `worker_index` of `num_workers`, such that each global step (across
all workers and replicas) will consume global_batch_size elements. The
returned value should be passed as the `batch_sizes` input parameter to
`tf.data.experimental.rebatch()`. The returned batch sizes meet the following
constraints:
Let G = global_batch_size, W = num_workers, R = num_replicas_per_worker
(A) for any worker, len(batch_sizes) = W * R
(B) for any worker, sum(batch_sizes) == G
(C) for any global step (i.e. R iterations on each worker), the sum of batches
consumed by replicas across all workers is G.
(D) any two batch sizes of any two replicas differs by at most one.
For example, suppose we have G = 7, W = 2, R = 2, and suppose we have two
files which each contain 7 elements:
```python
# WORKER 0
batch_sizes_0 = batch_sizes_for_worker(global_batch_size=global_batch_size,
num_workers=2,
num_replicas_per_worker=2,
worker_index=0)
print(batch_sizes_0)
>> [2, 2, 2, 1]
dataset_0 = tf.data.Dataset.from_tensor_slices(["file_a", "file_b"])
dataset_0 = dataset_0.shard(num_shards, index=0)
dataset_0 = dataset_0.batch(7)
dataset_0 = dataset_0.apply(tf.data.experimental.rebatch(batch_sizes_0))
for elem in dataset_0:
print(elem)
>> [[A0, A1], [A2, A3], [A4, A5], [A6]]
# WORKER 1
batch_sizes_1 = batch_sizes_for_worker(global_batch_size=global_batch_size,
num_workers=2,
num_replicas_per_worker=2,
worker_index=1)
print(batch_sizes_1)
>> [2, 1, 2, 2]
dataset_1 = tf.data.Dataset.from_tensor_slices(["file_a", "file_b"])
dataset_1 = dataset_1.shard(num_shards, index=1)
dataset_1 = dataset_1.batch(7)
dataset_1 = dataset_1.apply(tf.data.experimental.rebatch(batch_sizes_1))
for elem in dataset_1:
print(elem)
>> [[B0, B1], [B2], [B3, B4], [B5, B6]]
```
The above example will produce the following elements:
Step 1:
Worker 0 Replica 0: [A0, A1]
Worker 0 Replica 1: [A2, A3]
Worker 1 Replica 0: [B0, B1]
Worker 1 Replica 1: [B2]
Total batch size = 7
Step 2:
Worker 0 Replica 0: [A4, A5]
Worker 0 Replica 1: [A6]
Worker 1 Replica 0: [B3, B4]
Worker 1 Replica 1: [B5, B6]
Total batch size = 7
Args:
global_batch_size: A `tf.int64` scalar, representing the global batch size.
num_workers: An integer representing the number of workers the dataset will
be distributed across.
num_replicas_per_worker: An integer representing the number of replicas per
worker. All workers are assumed to have the same number of replicas.
worker_index: An integer index of the worker to be rebatched.
Returns:
A `tf.int64` vector, representing the batch sizes to rebatch the dataset
into.
"""
# Constraint (A)
num_subbatches = num_workers * num_replicas_per_worker
offset = worker_index * num_replicas_per_worker
const_value = tensor_util.constant_value(global_batch_size)
if const_value is not None:
# Use the constant global batch size for further calculations
global_batch_size = const_value
# Let N = W * R. Constraint (B) and (D) jointly mean that the iterations
# should have batch size either floor(B/N) or ceil(B/N). Namely, of the N
# subbatches a batch is split into, B - N * floor(B/N) of them will have size
# ceil(B/N), and the rest will have size floor(B/N).
floor = global_batch_size // num_subbatches
num_ceil = global_batch_size - (num_subbatches * floor)
# For worker 0, we assign the first num_ceil subbatches to have size
# ceil(B/N), and the remainder to have size floor(B/N). The other workers will
# each be offset by R * worker_index in order to meet constraint (C).
if const_value is not None:
# If the global batch size is a known constant value, we return a constant
# tensor directly instead of manipulating it with TF ops. This allows for
# better downstream shape inference.
worker_0 = [floor + 1] * num_ceil + [floor] * (num_subbatches - num_ceil)
return ops.convert_to_tensor(
worker_0[offset:] + worker_0[:offset],
dtype=dtypes.int64,
name="batch_sizes")
worker_0 = array_ops.ones(num_subbatches, dtype=dtypes.int64)
worker_0 = floor * worker_0 + array_ops.concat([
array_ops.ones(num_ceil, dtype=dtypes.int64),
array_ops.zeros(num_subbatches - num_ceil, dtype=dtypes.int64)
],
axis=0)
return array_ops.concat([worker_0[offset:], worker_0[:offset]], axis=0)
def compute_batch_size(dataset):
"""An operation that returns the batch size of the dataset.
This op tries to infer the batch size statically by walking up the dataset
tree from the final dataset node and returning the batch size of the first
batching dataset (such as from .batch() and .padded_batch()) that it
encounters. This differs from using the `element_spec` of a dataset in that it
does not account for partial batches.
This operation may fail if it encounters contradictory batch sizes (for
example, if the dataset is created by zipping together two datasets with
different batch sizes), if there are no explicit batching transformations, or
if there are operations downstream from the batching transformation that may
modify its batch size. In these cases, it returns a -1.
Args:
dataset: A `tf.data.Dataset` object.
Returns:
A `tf.int64` Tensor representing the batch size of the dataset sans partial
batches. If this cannot be inferred statically, the value of this tensor
will be -1.
"""
def get_static_batch_dim(type_spec):
try:
output_shape = type_spec._to_legacy_output_shapes() # pylint: disable=protected-access
except NotImplementedError:
return None
if not isinstance(output_shape, tensor_shape.TensorShape):
return None
if output_shape.rank is None:
return None
return output_shape.dims[0].value
batch_dims = [
get_static_batch_dim(type_spec)
for type_spec in nest.flatten(dataset_ops.get_structure(dataset))
]
if all(d is not None for d in batch_dims):
if all(d == batch_dims[0] for d in batch_dims):
# If all batch dimensions are known and equal, return that directly.
batch_dim = batch_dims[0]
else:
# If all batch dimensions are known but not all equal, return -1.
batch_dim = -1
return constant_op.constant(
batch_dim, dtype=dtypes.int64, name="static_batch_size")
# If any batch dimensions are unknown, use compute_batch_size op.
return ged_ops.compute_batch_size(dataset._variant_tensor) # pylint: disable=protected-access
_AutoShardDatasetV1.__doc__ = _AutoShardDataset.__doc__
tensorflow/python/data/experimental/ops/distributed_save_op.py
+111
View File
@@ -0,0 +1,111 @@
# Copyright 2022 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.
# ==============================================================================
"""Distributed saving of a dataset to disk."""
from typing import Optional
from tensorflow.core.protobuf import snapshot_pb2
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_experimental_dataset_ops
# TODO(b/238903802): Use TypeSpec serialization methods directly.
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.distributed_save")
def distributed_save(
dataset: dataset_ops.Dataset,
path: str,
data_service_address: str,
compression: str = "AUTO",
) -> Optional[ops.OperationType]:
"""Initiates the process of saving a dataset to disk using tf.data service.
The op uses tf.data service
(https://www.tensorflow.org/api_docs/python/tf/data/experimental/service) to
write a dataset snapshot. Returns immediately after submitting the request.
Does not wait for the snapshot to be finished. Requires that the tf.data
service run a fixed number of worker replicas.
To load the snapshot, users may optionally pass `wait=True` to
`tf.data.Dataset.load` so it can read snapshots as they are being written.
Example usage:
>>> import os
>>> import tempfile
>>> # Runs tf.data service.
>>> tempdir = tempfile.gettempdir()
>>> dispatcher = tf.data.experimental.service.DispatchServer(
... tf.data.experimental.service.DispatcherConfig(
... fault_tolerant_mode=True,
... work_dir=os.path.join(tempdir, "work_dir")))
>>> dispatcher_address = dispatcher.target.split("://")[1]
>>> worker = tf.data.experimental.service.WorkerServer(
... tf.data.experimental.service.WorkerConfig(
... dispatcher_address=dispatcher_address))
>>> # Writes dataset snapshot.
>>> path = os.path.join(tempdir, "dataset_snapshot")
>>> dataset = tf.data.Dataset.range(1)
>>> tf.data.experimental.distributed_save(dataset, path, dispatcher_address)
>>> # Loads a dataset snapshot.
>>> loaded_dataset = tf.data.Dataset.load(path, wait=True)
>>> for elem in loaded_dataset:
... print(elem)
tf.Tensor(0, shape=(), dtype=int64)
Args:
dataset: The `tf.data.Dataset` to save.
path: The directory path to save the dataset. Requires that:
- The directory does not exist and will create the directory.
- The file system supports atomic move (rename).
data_service_address: tf.data service dispatcher address.
compression: (Optional.) Whether and how to compress the `dataset` snapshot.
If `"AUTO"`, the tf.data runtime decides which algorithm to use. If
`"GZIP"` or `"SNAPPY"`, that specific algorithm is used. If `None`, the
`dataset` snapshot is not compressed.
Returns:
An operation which when executed performs the distributed save.
Raises:
ValueError: If `dispatcher_address` is invalid.
tf.errors.AlreadyExistsError: If the snapshot has already started or has
finished.
tf.errors.FailedPreconditionError: If the file system does not support
atomic move (rename).
tf.errors.InvalidArgumentError: If tf.data service is not running in the
fault tolerant mode.
"""
if not isinstance(data_service_address, str):
raise ValueError("`data_service_address` must be a string, but is a "
f"{type(data_service_address)} ({data_service_address}")
if not data_service_address:
raise ValueError("`data_service_address` must not be empty")
metadata = snapshot_pb2.DistributedSnapshotMetadata(
element_spec=nested_structure_coder.encode_structure(
dataset.element_spec).SerializeToString(),
compression=compression)
return gen_experimental_dataset_ops.distributed_save(
dataset._variant_tensor, # pylint: disable=protected-access
directory=path,
address=data_service_address,
metadata=metadata.SerializeToString())
tensorflow/python/data/experimental/ops/enumerate_ops.py
+54
View File
@@ -0,0 +1,54 @@
# Copyright 2017 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.
# ==============================================================================
"""Enumerate dataset transformations."""
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@deprecation.deprecated(None, "Use `tf.data.Dataset.enumerate()`.")
@tf_export("data.experimental.enumerate_dataset")
def enumerate_dataset(start=0):
"""A transformation that enumerates the elements of a dataset.
It is similar to python's `enumerate`.
For example:
```python
# NOTE: The following examples use `{ ... }` to represent the
# contents of a dataset.
a = { 1, 2, 3 }
b = { (7, 8), (9, 10) }
# The nested structure of the `datasets` argument determines the
# structure of elements in the resulting dataset.
a.apply(tf.data.experimental.enumerate_dataset(start=5))
=> { (5, 1), (6, 2), (7, 3) }
b.apply(tf.data.experimental.enumerate_dataset())
=> { (0, (7, 8)), (1, (9, 10)) }
```
Args:
start: A `tf.int64` scalar `tf.Tensor`, representing the start value for
enumeration.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.enumerate(start)
return _apply_fn
tensorflow/python/data/experimental/ops/error_ops.py
+51
View File
@@ -0,0 +1,51 @@
# Copyright 2017 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.
# ==============================================================================
"""Ignore_errors dataset transformations."""
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.ignore_errors")
@deprecation.deprecated(None, "Use `tf.data.Dataset.ignore_errors` instead.")
def ignore_errors(log_warning=False):
"""Creates a `Dataset` from another `Dataset` and silently ignores any errors.
Use this transformation to produce a dataset that contains the same elements
as the input, but silently drops any elements that caused an error. For
example:
```python
dataset = tf.data.Dataset.from_tensor_slices([1., 2., 0., 4.])
# Computing `tf.debugging.check_numerics(1. / 0.)` will raise an
InvalidArgumentError.
dataset = dataset.map(lambda x: tf.debugging.check_numerics(1. / x, "error"))
# Using `ignore_errors()` will drop the element that causes an error.
dataset =
dataset.apply(tf.data.experimental.ignore_errors()) # ==> {1., 0.5, 0.2}
```
Args:
log_warning: (Optional.) A 'tf.bool' scalar indicating whether ignored
errors should be logged to stderr. Defaults to 'False'.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.ignore_errors(log_warning)
return _apply_fn
tensorflow/python/data/experimental/ops/from_list.py
+122
View File
@@ -0,0 +1,122 @@
# Copyright 2022 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.
# ==============================================================================
"""Python API for creating a dataset from a list."""
import itertools
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest
from tensorflow.python.data.util import structure
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.util.tf_export import tf_export
class _ListDataset(dataset_ops.DatasetSource):
"""A `Dataset` of elements from a list."""
def __init__(self, elements, name=None):
if not elements:
raise ValueError(
"Invalid `elements`. `elements` should not be empty. If you want an"
" empty dataset, use `tf.data.Dataset.range(0)`."
)
if not isinstance(elements, list):
raise ValueError("Invalid `elements`. `elements` must be a list.")
elements = [structure.normalize_element(element) for element in elements]
type_specs = [
structure.type_spec_from_value(element) for element in elements
]
# Check that elements have same nested structure.
num_elements = len(elements)
for i in range(1, num_elements):
nest.assert_same_structure(type_specs[0], type_specs[i])
# Infer elements' supershape.
flattened_type_specs = [nest.flatten(type_spec) for type_spec in type_specs]
num_tensors_per_element = len(flattened_type_specs[0])
flattened_structure = [None] * num_tensors_per_element
for i in range(num_tensors_per_element):
flattened_structure[i] = flattened_type_specs[0][i]
for j in range(1, num_elements):
flattened_structure[i] = flattened_structure[
i].most_specific_common_supertype([flattened_type_specs[j][i]])
if not isinstance(type_specs[0], dataset_ops.DatasetSpec):
self._tensors = list(
itertools.chain.from_iterable(
[nest.flatten(element) for element in elements]))
else:
self._tensors = [x._variant_tensor for x in elements]
self._structure = nest.pack_sequence_as(type_specs[0], flattened_structure)
self._name = name
variant_tensor = gen_experimental_dataset_ops.list_dataset(
self._tensors,
output_types=self._flat_types,
output_shapes=self._flat_shapes,
metadata=self._metadata.SerializeToString())
super(_ListDataset, self).__init__(variant_tensor)
@property
def element_spec(self):
return self._structure
@tf_export("data.experimental.from_list")
def from_list(elements, name=None):
"""Creates a `Dataset` comprising the given list of elements.
The returned dataset will produce the items in the list one by one. The
functionality is identical to `Dataset.from_tensor_slices` when elements are
scalars, but different when elements have structure. Consider the following
example.
>>> dataset = tf.data.experimental.from_list([(1, 'a'), (2, 'b'), (3, 'c')])
>>> [(n.item(), s) for n, s in dataset.as_numpy_iterator()]
[(1, b'a'), (2, b'b'), (3, b'c')]
To get the same output with `from_tensor_slices`, the data needs to be
reorganized:
>>> dataset = tf.data.Dataset.from_tensor_slices(([1, 2, 3], ['a', 'b', 'c']))
>>> [(n.item(), s) for n, s in dataset.as_numpy_iterator()]
[(1, b'a'), (2, b'b'), (3, b'c')]
Unlike `from_tensor_slices`, `from_list` supports non-rectangular input:
>>> dataset = tf.data.experimental.from_list([[1], [2, 3]])
>>> list(dataset.as_numpy_iterator())
[array([1], dtype=int32), array([2, 3], dtype=int32)]
Achieving the same with `from_tensor_slices` requires the use of ragged
tensors.
`from_list` can be more performant than `from_tensor_slices` in some cases,
since it avoids the need for data slicing each epoch. However, it can also be
less performant, because data is stored as many small tensors rather than a
few large tensors as in `from_tensor_slices`. The general guidance is to
prefer `from_list` from a performance perspective when the number of elements
is small (less than 1000).
Args:
elements: A list of elements whose components have the same nested
structure.
name: (Optional.) A name for the tf.data operation.
Returns:
Dataset: A `Dataset` of the `elements`.
"""
return _ListDataset(elements, name)
tensorflow/python/data/experimental/ops/get_single_element.py
+109
View File
@@ -0,0 +1,109 @@
# Copyright 2017 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.
# ==============================================================================
"""Python wrappers for Datasets and Iterators."""
from tensorflow.python.types import data as data_types
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@deprecation.deprecated(None, "Use `tf.data.Dataset.get_single_element()`.")
@tf_export("data.experimental.get_single_element")
def get_single_element(dataset):
"""Returns the single element of the `dataset` as a nested structure of tensors.
The function enables you to use a `tf.data.Dataset` in a stateless
"tensor-in tensor-out" expression, without creating an iterator.
This facilitates the ease of data transformation on tensors using the
optimized `tf.data.Dataset` abstraction on top of them.
For example, lets consider a `preprocessing_fn` which would take as an
input the raw features and returns the processed feature along with
it's label.
```python
def preprocessing_fn(raw_feature):
# ... the raw_feature is preprocessed as per the use-case
return feature
raw_features = ... # input batch of BATCH_SIZE elements.
dataset = (tf.data.Dataset.from_tensor_slices(raw_features)
.map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
.batch(BATCH_SIZE))
processed_features = tf.data.experimental.get_single_element(dataset)
```
In the above example, the `raw_features` tensor of length=BATCH_SIZE
was converted to a `tf.data.Dataset`. Next, each of the `raw_feature` was
mapped using the `preprocessing_fn` and the processed features were
grouped into a single batch. The final `dataset` contains only one element
which is a batch of all the processed features.
NOTE: The `dataset` should contain only one element.
Now, instead of creating an iterator for the `dataset` and retrieving the
batch of features, the `tf.data.experimental.get_single_element()` function
is used to skip the iterator creation process and directly output the batch
of features.
This can be particularly useful when your tensor transformations are
expressed as `tf.data.Dataset` operations, and you want to use those
transformations while serving your model.
# Keras
```python
model = ... # A pre-built or custom model
class PreprocessingModel(tf.keras.Model):
def __init__(self, model):
super().__init__(self)
self.model = model
@tf.function(input_signature=[...])
def serving_fn(self, data):
ds = tf.data.Dataset.from_tensor_slices(data)
ds = ds.map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
ds = ds.batch(batch_size=BATCH_SIZE)
return tf.argmax(
self.model(tf.data.experimental.get_single_element(ds)),
axis=-1
)
preprocessing_model = PreprocessingModel(model)
your_exported_model_dir = ... # save the model to this path.
tf.saved_model.save(preprocessing_model, your_exported_model_dir,
signatures={'serving_default': preprocessing_model.serving_fn})
```
Args:
dataset: A `tf.data.Dataset` object containing a single element.
Returns:
A nested structure of `tf.Tensor` objects, corresponding to the single
element of `dataset`.
Raises:
TypeError: if `dataset` is not a `tf.data.Dataset` object.
InvalidArgumentError: (at runtime) if `dataset` does not contain exactly
one element.
"""
if not isinstance(dataset, data_types.DatasetV2):
raise TypeError(
f"Invalid `dataset`. Expected a `tf.data.Dataset` object "
f"but got {type(dataset)}.")
return dataset.get_single_element()
tensorflow/python/data/experimental/ops/global_shuffle_op.py
+98
View File
@@ -0,0 +1,98 @@
# Copyright 2024 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.
# ==============================================================================
"""Globally shuffles tf.data datasets."""
from typing import Optional, Union
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import options as options_lib
from tensorflow.python.data.util import random_seed
from tensorflow.python.framework import tensor
from tensorflow.python.ops import gen_dataset_ops
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
def _global_shuffle( # pylint: disable=unused-private-name
input_dataset: dataset_ops.DatasetV2,
seed: Optional[Union[int, tensor.Tensor]] = None,
reshuffle_each_iteration: bool = True,
name: Optional[str] = None) -> dataset_ops.DatasetV2:
"""Globally shuffles the elements of `input_dataset`.
The shuffling is done efficiently, without needing to buffer any additional
data. To achieve this, the transformations preceding global_shuffle must all
support random access.
Requires that:
- The shuffled dataset and all its input datasets support random access.
- The input_dataset to have a known, finite cardinality. Users can use
`tf.data.experimental.assert_cardinality` to specify the cardinality of a
dataset if it cannot be determined at runtime.
TODO(b/325112575): Move the API to dataset_ops.py.
Args:
input_dataset: The dataset to be shuffled.
seed: An int or `tf.int64` scalar `tf.Tensor` to control the shuffle order.
If `None`, a random seed will be used.
reshuffle_each_iteration: A boolean, which if True, indicates that a
different shuffle order should be generated for each iteration of the
dataset. (Defaults to `True`.)
name: (Optional.) A name for the tf.data operation.
Returns:
A new `Dataset` where elements are produced in a globally shuffled order.
Raises:
- InvalidArgumentError if the input dataset does not support random access,
or it has infinite or unknown cardinality.
- FailedPreconditionError for batching with `drop_remainder=False`.
"""
return _GlobalShuffleDataset(
input_dataset,
seed=seed,
reshuffle_each_iteration=reshuffle_each_iteration,
name=name)
class _GlobalShuffleDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""Shuffles all elements in the input dataset."""
def __init__(
self,
input_dataset: dataset_ops.DatasetV2,
seed: Optional[Union[int, tensor.Tensor]] = None,
reshuffle_each_iteration: bool = True,
name: Optional[str] = None):
options = options_lib.Options()
# Currently, prefetching threads cannot access the runtime context required
# for global shuffling when `warm_start` is enabled. Supporting it will be
# future work.
options.experimental_warm_start = False
input_dataset = input_dataset.with_options(options)
self._input_dataset = input_dataset
self._seed, self._seed2 = random_seed.get_seed(seed)
self._reshuffle_each_iteration = reshuffle_each_iteration
self._name = name
variant_tensor = ged_ops.global_shuffle_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
seed=self._seed,
seed2=self._seed2,
seed_generator=gen_dataset_ops.dummy_seed_generator(),
reshuffle_each_iteration=self._reshuffle_each_iteration,
**self._common_args)
super().__init__(input_dataset, variant_tensor)
tensorflow/python/data/experimental/ops/grouping.py
+428
View File
@@ -0,0 +1,428 @@
# Copyright 2017 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.
# ==============================================================================
"""Grouping dataset transformations."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import structured_function
from tensorflow.python.data.util import nest
from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.group_by_reducer")
def group_by_reducer(key_func, reducer):
"""A transformation that groups elements and performs a reduction.
This transformation maps element of a dataset to a key using `key_func` and
groups the elements by key. The `reducer` is used to process each group; its
`init_func` is used to initialize state for each group when it is created, the
`reduce_func` is used to update the state every time an element is mapped to
the matching group, and the `finalize_func` is used to map the final state to
an output value.
Args:
key_func: A function mapping a nested structure of tensors
(having shapes and types defined by `self.output_shapes` and
`self.output_types`) to a scalar `tf.int64` tensor.
reducer: An instance of `Reducer`, which captures the reduction logic using
the `init_func`, `reduce_func`, and `finalize_func` functions.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
return _GroupByReducerDataset(dataset, key_func, reducer)
return _apply_fn
@deprecation.deprecated(None, "Use `tf.data.Dataset.group_by_window(...)`.")
@tf_export("data.experimental.group_by_window")
def group_by_window(key_func,
reduce_func,
window_size=None,
window_size_func=None):
"""A transformation that groups windows of elements by key and reduces them.
This transformation maps each consecutive element in a dataset to a key
using `key_func` and groups the elements by key. It then applies
`reduce_func` to at most `window_size_func(key)` elements matching the same
key. All except the final window for each key will contain
`window_size_func(key)` elements; the final window may be smaller.
You may provide either a constant `window_size` or a window size determined by
the key through `window_size_func`.
Args:
key_func: A function mapping a nested structure of tensors
(having shapes and types defined by `self.output_shapes` and
`self.output_types`) to a scalar `tf.int64` tensor.
reduce_func: A function mapping a key and a dataset of up to `window_size`
consecutive elements matching that key to another dataset.
window_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
consecutive elements matching the same key to combine in a single
batch, which will be passed to `reduce_func`. Mutually exclusive with
`window_size_func`.
window_size_func: A function mapping a key to a `tf.int64` scalar
`tf.Tensor`, representing the number of consecutive elements matching
the same key to combine in a single batch, which will be passed to
`reduce_func`. Mutually exclusive with `window_size`.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
Raises:
ValueError: if neither or both of {`window_size`, `window_size_func`} are
passed.
"""
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
return dataset.group_by_window(
key_func=key_func,
reduce_func=reduce_func,
window_size=window_size,
window_size_func=window_size_func)
return _apply_fn
@deprecation.deprecated(None,
"Use `tf.data.Dataset.bucket_by_sequence_length(...)`.")
@tf_export("data.experimental.bucket_by_sequence_length")
def bucket_by_sequence_length(element_length_func,
bucket_boundaries,
bucket_batch_sizes,
padded_shapes=None,
padding_values=None,
pad_to_bucket_boundary=False,
no_padding=False,
drop_remainder=False):
"""A transformation that buckets elements in a `Dataset` by length.
Elements of the `Dataset` are grouped together by length and then are padded
and batched.
This is useful for sequence tasks in which the elements have variable length.
Grouping together elements that have similar lengths reduces the total
fraction of padding in a batch which increases training step efficiency.
Below is an example to bucketize the input data to the 3 buckets
"[0, 3), [3, 5), [5, inf)" based on sequence length, with batch size 2.
>>> elements = [
... [0], [1, 2, 3, 4], [5, 6, 7],
... [7, 8, 9, 10, 11], [13, 14, 15, 16, 19, 20], [21, 22]]
>>> dataset = tf.data.Dataset.from_generator(
... lambda: elements, tf.int64, output_shapes=[None])
>>> dataset = dataset.apply(
... tf.data.experimental.bucket_by_sequence_length(
... element_length_func=lambda elem: tf.shape(elem)[0],
... bucket_boundaries=[3, 5],
... bucket_batch_sizes=[2, 2, 2]))
>>> for elem in dataset.as_numpy_iterator():
... print(elem)
[[1 2 3 4]
[5 6 7 0]]
[[ 7 8 9 10 11 0]
[13 14 15 16 19 20]]
[[ 0 0]
[21 22]]
There is also a possibility to pad the dataset till the bucket boundary.
You can also provide which value to be used while padding the data.
Below example uses `-1` as padding and it also shows the input data
being bucketizied to two buckets "[0,3], [4,6]".
>>> elements = [
... [0], [1, 2, 3, 4], [5, 6, 7],
... [7, 8, 9, 10, 11], [13, 14, 15, 16, 19, 20], [21, 22]]
>>> dataset = tf.data.Dataset.from_generator(
... lambda: elements, tf.int32, output_shapes=[None])
>>> dataset = dataset.apply(
... tf.data.experimental.bucket_by_sequence_length(
... element_length_func=lambda elem: tf.shape(elem)[0],
... bucket_boundaries=[4, 7],
... bucket_batch_sizes=[2, 2, 2],
... pad_to_bucket_boundary=True,
... padding_values=-1))
>>> for elem in dataset.as_numpy_iterator():
... print(elem)
[[ 0 -1 -1]
[ 5 6 7]]
[[ 1 2 3 4 -1 -1]
[ 7 8 9 10 11 -1]]
[[21 22 -1]]
[[13 14 15 16 19 20]]
When using `pad_to_bucket_boundary` option, it can be seen that it is
not always possible to maintain the bucket batch size.
You can drop the batches that do not maintain the bucket batch size by
using the option `drop_remainder`. Using the same input data as in the
above example you get the following result.
>>> elements = [
... [0], [1, 2, 3, 4], [5, 6, 7],
... [7, 8, 9, 10, 11], [13, 14, 15, 16, 19, 20], [21, 22]]
>>> dataset = tf.data.Dataset.from_generator(
... lambda: elements, tf.int32, output_shapes=[None])
>>> dataset = dataset.apply(
... tf.data.experimental.bucket_by_sequence_length(
... element_length_func=lambda elem: tf.shape(elem)[0],
... bucket_boundaries=[4, 7],
... bucket_batch_sizes=[2, 2, 2],
... pad_to_bucket_boundary=True,
... padding_values=-1,
... drop_remainder=True))
>>> for elem in dataset.as_numpy_iterator():
... print(elem)
[[ 0 -1 -1]
[ 5 6 7]]
[[ 1 2 3 4 -1 -1]
[ 7 8 9 10 11 -1]]
Args:
element_length_func: function from element in `Dataset` to `tf.int32`,
determines the length of the element, which will determine the bucket it
goes into.
bucket_boundaries: `list<int>`, upper length boundaries of the buckets.
bucket_batch_sizes: `list<int>`, batch size per bucket. Length should be
`len(bucket_boundaries) + 1`.
padded_shapes: Nested structure of `tf.TensorShape` to pass to
`tf.data.Dataset.padded_batch`. If not provided, will use
`dataset.output_shapes`, which will result in variable length dimensions
being padded out to the maximum length in each batch.
padding_values: Values to pad with, passed to
`tf.data.Dataset.padded_batch`. Defaults to padding with 0.
pad_to_bucket_boundary: bool, if `False`, will pad dimensions with unknown
size to maximum length in batch. If `True`, will pad dimensions with
unknown size to bucket boundary minus 1 (i.e., the maximum length in each
bucket), and caller must ensure that the source `Dataset` does not contain
any elements with length longer than `max(bucket_boundaries)`.
no_padding: `bool`, indicates whether to pad the batch features (features
need to be either of type `tf.sparse.SparseTensor` or of same shape).
drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
whether the last batch should be dropped in the case it has fewer than
`batch_size` elements; the default behavior is not to drop the smaller
batch.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
Raises:
ValueError: if `len(bucket_batch_sizes) != len(bucket_boundaries) + 1`.
"""
def _apply_fn(dataset):
return dataset.bucket_by_sequence_length(
element_length_func=element_length_func,
bucket_boundaries=bucket_boundaries,
bucket_batch_sizes=bucket_batch_sizes,
padded_shapes=padded_shapes,
padding_values=padding_values,
pad_to_bucket_boundary=pad_to_bucket_boundary,
no_padding=no_padding,
drop_remainder=drop_remainder)
return _apply_fn
class _GroupByReducerDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that groups its input and performs a reduction."""
def __init__(self, input_dataset, key_func, reducer):
"""See `group_by_reducer()` for details."""
self._input_dataset = input_dataset
self._make_key_func(key_func, input_dataset)
self._make_init_func(reducer.init_func)
self._make_reduce_func(reducer.reduce_func, input_dataset)
self._make_finalize_func(reducer.finalize_func)
variant_tensor = ged_ops.experimental_group_by_reducer_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._key_func.function.captured_inputs,
self._init_func.function.captured_inputs,
self._reduce_func.function.captured_inputs,
self._finalize_func.function.captured_inputs,
key_func=self._key_func.function,
init_func=self._init_func.function,
reduce_func=self._reduce_func.function,
finalize_func=self._finalize_func.function,
**self._flat_structure)
super(_GroupByReducerDataset, self).__init__(input_dataset, variant_tensor)
def _make_key_func(self, key_func, input_dataset):
"""Make wrapping defun for key_func."""
self._key_func = structured_function.StructuredFunctionWrapper(
key_func, self._transformation_name(), dataset=input_dataset)
if not self._key_func.output_structure.is_compatible_with(
tensor_spec.TensorSpec([], dtypes.int64)):
raise ValueError(
f"Invalid `key_func`. Expected `key_func` to return a scalar "
f"tf.int64 tensor, but instead `key_func` has output "
f"types={self._key_func.output_types} "
f"and shapes={self._key_func.output_shapes}."
)
def _make_init_func(self, init_func):
"""Make wrapping defun for init_func."""
self._init_func = structured_function.StructuredFunctionWrapper(
init_func,
self._transformation_name(),
input_structure=tensor_spec.TensorSpec([], dtypes.int64))
def _make_reduce_func(self, reduce_func, input_dataset):
"""Make wrapping defun for reduce_func."""
# Iteratively rerun the reduce function until reaching a fixed point on
# `self._state_structure`.
self._state_structure = self._init_func.output_structure
state_types = self._init_func.output_types
state_shapes = self._init_func.output_shapes
state_classes = self._init_func.output_classes
need_to_rerun = True
while need_to_rerun:
wrapped_func = structured_function.StructuredFunctionWrapper(
reduce_func,
self._transformation_name(),
input_structure=(self._state_structure, input_dataset.element_spec),
add_to_graph=False)
# Extract and validate class information from the returned values.
for new_state_class, state_class in zip(
nest.flatten(wrapped_func.output_classes),
nest.flatten(state_classes)):
if not issubclass(new_state_class, state_class):
raise TypeError(
f"Invalid `reducer`. The output class of the "
f"`reducer.reduce_func` {wrapped_func.output_classes}, "
f"does not match the class of the reduce state "
f"{self._state_classes}.")
# Extract and validate type information from the returned values.
for new_state_type, state_type in zip(
nest.flatten(wrapped_func.output_types), nest.flatten(state_types)):
if new_state_type != state_type:
raise TypeError(
f"Invalid `reducer`. The element types for the new state "
f"{wrapped_func.output_types} do not match the element types "
f"of the old state {self._init_func.output_types}."
)
# Extract shape information from the returned values.
flat_state_shapes = nest.flatten(state_shapes)
flat_new_state_shapes = nest.flatten(wrapped_func.output_shapes)
weakened_state_shapes = [
original.most_specific_compatible_shape(new)
for original, new in zip(flat_state_shapes, flat_new_state_shapes)
]
need_to_rerun = False
for original_shape, weakened_shape in zip(flat_state_shapes,
weakened_state_shapes):
if original_shape.ndims is not None and (
weakened_shape.ndims is None or
original_shape.as_list() != weakened_shape.as_list()):
need_to_rerun = True
break
if need_to_rerun:
state_shapes = nest.pack_sequence_as(
self._init_func.output_shapes, weakened_state_shapes)
self._state_structure = structure.convert_legacy_structure(
state_types, state_shapes, state_classes)
self._reduce_func = wrapped_func
self._reduce_func.function.add_to_graph(ops.get_default_graph())
def _make_finalize_func(self, finalize_func):
"""Make wrapping defun for finalize_func."""
self._finalize_func = structured_function.StructuredFunctionWrapper(
finalize_func,
self._transformation_name(),
input_structure=self._state_structure)
@property
def element_spec(self):
return self._finalize_func.output_structure
def _functions(self):
return [
self._key_func, self._init_func, self._reduce_func, self._finalize_func
]
def _transformation_name(self):
return "tf.data.experimental.group_by_reducer()"
@tf_export("data.experimental.Reducer")
class Reducer:
"""A reducer is used for reducing a set of elements.
A reducer is represented as a tuple of the three functions:
- init_func - to define initial value: key => initial state
- reducer_func - operation to perform on values with same key: (old state, input) => new state
- finalize_func - value to return in the end: state => result
For example,
```
def init_func(_):
return (0.0, 0.0)
def reduce_func(state, value):
return (state[0] + value['features'], state[1] + 1)
def finalize_func(s, n):
return s / n
reducer = tf.data.experimental.Reducer(init_func, reduce_func, finalize_func)
```
"""
def __init__(self, init_func, reduce_func, finalize_func):
self._init_func = init_func
self._reduce_func = reduce_func
self._finalize_func = finalize_func
@property
def init_func(self):
return self._init_func
@property
def reduce_func(self):
return self._reduce_func
@property
def finalize_func(self):
return self._finalize_func
tensorflow/python/data/experimental/ops/index_flat_map_op.py
+144
View File
@@ -0,0 +1,144 @@
# Copyright 2024 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.
# ==============================================================================
"""Python API for `index_flat_map` dataset, which supports global shuffling."""
from typing import Any, Callable, Optional, Sequence, Union
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import structured_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
_IndexType = Union[int, tensor.Tensor]
def index_flat_map( # pylint: disable=unused-private-name
input_dataset: dataset_ops.Dataset,
map_func: Callable[[Any], tensor.Tensor],
index_map_func: Callable[[_IndexType], tuple[_IndexType, _IndexType]],
*,
output_cardinality: _IndexType = dataset_ops.UNKNOWN,
name: Optional[str] = None) -> dataset_ops.Dataset:
"""A variant of flat_map that supports global shuffling.
In addition to a `map_func`, the user needs to provide an `index_map_fn`.
Given an index in the flattened dataset, the `index_map_fn` returns an
(element_index, offset) tuple that represents the index of the element in the
unflattened dataset, and the offset in the unflattened element.
For example, users could flatten a dataset as following:
def _split(element: str) -> tensor.Tensor:
return tf.strings.split(element, " ")
def _index_map_func(flattened_index: int) -> tuple[int, int]:
'''Returns an (element_index, offset) tuple for the requested element index.
For a `flattened_index` that represents an element index in the output
dataset, this function returns a tuple that represents the index of the
element in the unflattened dataset, and the offset in that element.
See the example below on how to implement this function.
'''
dataset = tf.data.Dataset.from_tensor_slices([["0 1", "2 3 4 5", "6 7", "8"])
dataset = index_flat_map_op.index_flat_map(dataset, _split, _index_map_func)
for x in dataset:
print(x) # Produces "0", "1", "2", "3", "4", "5", "6", "7", "8".
Given an input of 5, the `index_map_func` should return (1, 3), which means
that the target element in the result dataset is in the 2nd element in the
original dataset ("2 3 4 5") with an offset of 3, which is "5".
Args:
input_dataset: The input dataset.
map_func: A function mapping a dataset element to a Tensor of the mapped
elements.
index_map_func: Given an index in the flattened dataset, returns a
(element index, offset) tuple that represents the index of the element in
the unflattened dataset, and the offset in the unflattened element.
output_cardinality: Cardinality of the output dataset. Can be an int or a
Tensor of int64. Required if the dataset is globally shuffled because the
cardinality cannot be inferred as `map_func` produces a varied number of
output elements from each input element.
name: (Optional.) A name for the tf.data operation.
Returns:
A new `Dataset` with the transformation applied as described above.
Raises:
errors.InvalidArgumentError: If `index_map_func` does not return a tuple of
two integers, or the returned offset is out of range.
"""
return _IndexFlatMapDataset(
input_dataset, map_func, index_map_func, output_cardinality, name)
class _IndexFlatMapDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that maps a function over its input and flattens the result."""
def __init__(
self,
input_dataset: dataset_ops.Dataset,
map_func: Callable[[Any], tensor.Tensor],
index_map_func: Callable[[_IndexType], tuple[_IndexType, _IndexType]],
output_cardinality: _IndexType = dataset_ops.UNKNOWN,
name: str = None):
self._input_dataset = input_dataset
self._map_func = structured_function.StructuredFunctionWrapper(
map_func,
transformation_name=f"{self._transformation_name()}.map_func",
dataset=input_dataset)
self._index_map_func = structured_function.StructuredFunctionWrapper(
index_map_func,
transformation_name=f"{self._transformation_name()}.index_map_func",
input_structure=tensor_spec.TensorSpec([], dtypes.int64))
self._output_cardinality = ops.convert_to_tensor(
output_cardinality, dtype=dtypes.int64)
self._name = name
variant_tensor = ged_ops.index_flat_map_dataset(
input_dataset._variant_tensor,
self._map_func.function.captured_inputs,
self._index_map_func.function.captured_inputs,
map_func=self._map_func.function,
index_map_func=self._index_map_func.function,
output_cardinality=self._output_cardinality,
**self._common_args)
super().__init__(input_dataset, variant_tensor)
@property
def element_spec(self) -> Any:
output_structure = self._map_func.output_structure
# If the `map_func` returns a Python list of lists, then `output_structure`
# contains s sequence of `TensorSpecs`, each representing the structure of
# the inner lists.
if isinstance(output_structure, Sequence):
return output_structure[0]
# If the `map_func` returns a Tensor of nested lists, then each mapped
# element is one stacked Tensor.
if output_structure.shape.rank > 1:
return tensor_spec.TensorSpec(
shape=output_structure.shape[1:], dtype=output_structure.dtype)
# `map_func` returns a list of scalars.
return tensor_spec.TensorSpec(shape=[], dtype=output_structure.dtype)
def _functions(self) -> list[structured_function.StructuredFunctionWrapper]:
return [self._map_func, self._index_map_func]
def _transformation_name(self) -> str:
return "Dataset.index_flat_map()"
tensorflow/python/data/experimental/ops/interleave_ops.py
+261
View File
@@ -0,0 +1,261 @@
# Copyright 2017 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.
# ==============================================================================
"""Non-deterministic dataset transformations."""
from tensorflow.python import tf2
from tensorflow.python.compat import v2_compat
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import readers
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@deprecation.deprecated(
None,
"Use `tf.data.Dataset.interleave(map_func, cycle_length, block_length, "
"num_parallel_calls=tf.data.AUTOTUNE)` instead. If sloppy "
"execution is desired, use `tf.data.Options.deterministic`.")
@tf_export("data.experimental.parallel_interleave")
def parallel_interleave(map_func,
cycle_length,
block_length=1,
sloppy=False,
buffer_output_elements=None,
prefetch_input_elements=None):
"""A parallel version of the `Dataset.interleave()` transformation.
`parallel_interleave()` maps `map_func` across its input to produce nested
datasets, and outputs their elements interleaved. Unlike
`tf.data.Dataset.interleave`, it gets elements from `cycle_length` nested
datasets in parallel, which increases the throughput, especially in the
presence of stragglers. Furthermore, the `sloppy` argument can be used to
improve performance, by relaxing the requirement that the outputs are produced
in a deterministic order, and allowing the implementation to skip over nested
datasets whose elements are not readily available when requested.
Example usage:
```python
# Preprocess 4 files concurrently.
filenames = tf.data.Dataset.list_files("/path/to/data/train*.tfrecords")
dataset = filenames.apply(
tf.data.experimental.parallel_interleave(
lambda filename: tf.data.TFRecordDataset(filename),
cycle_length=4))
```
WARNING: If `sloppy` is `True`, the order of produced elements is not
deterministic.
Args:
map_func: A function mapping a nested structure of tensors to a `Dataset`.
cycle_length: The number of input `Dataset`s to interleave from in parallel.
block_length: The number of consecutive elements to pull from an input
`Dataset` before advancing to the next input `Dataset`.
sloppy: A boolean controlling whether determinism should be traded for
performance by allowing elements to be produced out of order. If `sloppy`
is `None`, the `tf.data.Options.deterministic` dataset option (`True` by
default) is used to decide whether to enforce a deterministic order.
buffer_output_elements: The number of elements each iterator being
interleaved should buffer (similar to the `.prefetch()` transformation for
each interleaved iterator).
prefetch_input_elements: The number of input elements to transform to
iterators before they are needed for interleaving.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return readers.ParallelInterleaveDataset(dataset, map_func, cycle_length,
block_length, sloppy,
buffer_output_elements,
prefetch_input_elements)
return _apply_fn
@deprecation.deprecated(None,
"Use `tf.data.Dataset.sample_from_datasets(...)`.")
@tf_export("data.experimental.sample_from_datasets", v1=[])
def sample_from_datasets_v2(datasets,
weights=None,
seed=None,
stop_on_empty_dataset=False):
"""Samples elements at random from the datasets in `datasets`.
Creates a dataset by interleaving elements of `datasets` with `weight[i]`
probability of picking an element from dataset `i`. Sampling is done without
replacement. For example, suppose we have 2 datasets:
```python
dataset1 = tf.data.Dataset.range(0, 3)
dataset2 = tf.data.Dataset.range(100, 103)
```
Suppose also that we sample from these 2 datasets with the following weights:
```python
sample_dataset = tf.data.Dataset.sample_from_datasets(
[dataset1, dataset2], weights=[0.5, 0.5])
```
One possible outcome of elements in sample_dataset is:
```
print(list(sample_dataset.as_numpy_iterator()))
# [100, 0, 1, 101, 2, 102]
```
Args:
datasets: A non-empty list of `tf.data.Dataset` objects with compatible
structure.
weights: (Optional.) A list or Tensor of `len(datasets)` floating-point
values where `weights[i]` represents the probability to sample from
`datasets[i]`, or a `tf.data.Dataset` object where each element is such a
list. Defaults to a uniform distribution across `datasets`.
seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the random
seed that will be used to create the distribution. See
`tf.random.set_seed` for behavior.
stop_on_empty_dataset: If `True`, sampling stops if it encounters an empty
dataset. If `False`, it skips empty datasets. It is recommended to set it
to `True`. Otherwise, the distribution of samples starts off as the user
intends, but may change as input datasets become empty. This can be
difficult to detect since the dataset starts off looking correct. Default
to `False` for backward compatibility.
Returns:
A dataset that interleaves elements from `datasets` at random, according to
`weights` if provided, otherwise with uniform probability.
Raises:
TypeError: If the `datasets` or `weights` arguments have the wrong type.
ValueError:
- If `datasets` is empty, or
- If `weights` is specified and does not match the length of `datasets`.
"""
return dataset_ops.Dataset.sample_from_datasets(
datasets=datasets,
weights=weights,
seed=seed,
stop_on_empty_dataset=stop_on_empty_dataset)
@deprecation.deprecated(None,
"Use `tf.data.Dataset.sample_from_datasets(...)`.")
@tf_export(v1=["data.experimental.sample_from_datasets"])
def sample_from_datasets_v1(datasets,
weights=None,
seed=None,
stop_on_empty_dataset=False):
return dataset_ops.DatasetV1Adapter(
sample_from_datasets_v2(datasets, weights, seed, stop_on_empty_dataset))
sample_from_datasets_v1.__doc__ = sample_from_datasets_v2.__doc__
@deprecation.deprecated(
None, "Use `tf.data.Dataset.choose_from_datasets(...)` instead. Note that, "
"unlike the experimental endpoint, the non-experimental endpoint "
"sets `stop_on_empty_dataset=True` by default. You should set this "
"argument explicitly in case you would like to match the behavior of the "
"experimental endpoint.")
@tf_export("data.experimental.choose_from_datasets", v1=[])
def choose_from_datasets_v2(datasets,
choice_dataset,
stop_on_empty_dataset=False):
"""Creates a dataset that deterministically chooses elements from `datasets`.
For example, given the following datasets:
```python
datasets = [tf.data.Dataset.from_tensors("foo").repeat(),
tf.data.Dataset.from_tensors("bar").repeat(),
tf.data.Dataset.from_tensors("baz").repeat()]
# Define a dataset containing `[0, 1, 2, 0, 1, 2, 0, 1, 2]`.
choice_dataset = tf.data.Dataset.range(3).repeat(3)
result = tf.data.experimental.choose_from_datasets(datasets, choice_dataset)
```
The elements of `result` will be:
```
"foo", "bar", "baz", "foo", "bar", "baz", "foo", "bar", "baz"
```
Args:
datasets: A non-empty list of `tf.data.Dataset` objects with compatible
structure.
choice_dataset: A `tf.data.Dataset` of scalar `tf.int64` tensors between `0`
and `len(datasets) - 1`.
stop_on_empty_dataset: If `True`, selection stops if it encounters an empty
dataset. If `False`, it skips empty datasets. It is recommended to set it
to `True`. Otherwise, the selected elements start off as the user intends,
but may change as input datasets become empty. This can be difficult to
detect since the dataset starts off looking correct. Default to `False`
for backward compatibility.
Returns:
A dataset that interleaves elements from `datasets` according to the values
of `choice_dataset`.
Raises:
TypeError: If `datasets` or `choice_dataset` has the wrong type.
ValueError: If `datasets` is empty.
"""
return dataset_ops.Dataset.choose_from_datasets(
datasets=datasets,
choice_dataset=choice_dataset,
stop_on_empty_dataset=stop_on_empty_dataset)
@deprecation.deprecated(
None, "Use `tf.data.Dataset.choose_from_datasets(...)` instead. Note that, "
"unlike the experimental endpoint, the non-experimental endpoint "
"sets `stop_on_empty_dataset=True` by default. You should set this "
"argument explicitly in case you would like to match the behavior of the "
"experimental endpoint.")
@tf_export(v1=["data.experimental.choose_from_datasets"])
def choose_from_datasets_v1(datasets,
choice_dataset,
stop_on_empty_dataset=False):
return dataset_ops.DatasetV1Adapter(
choose_from_datasets_v2(datasets, choice_dataset, stop_on_empty_dataset))
choose_from_datasets_v1.__doc__ = choose_from_datasets_v2.__doc__
if tf2.enabled():
choose_from_datasets = choose_from_datasets_v2
sample_from_datasets = sample_from_datasets_v2
else:
choose_from_datasets = choose_from_datasets_v1
sample_from_datasets = sample_from_datasets_v1
def _tf2_callback():
global choose_from_datasets, sample_from_datasets
if tf2.enabled():
choose_from_datasets = choose_from_datasets_v2
sample_from_datasets = sample_from_datasets_v2
else:
choose_from_datasets = choose_from_datasets_v1
sample_from_datasets = sample_from_datasets_v1
v2_compat.register_data_v2_callback(_tf2_callback)
tensorflow/python/data/experimental/ops/io.py
+166
View File
@@ -0,0 +1,166 @@
# Copyright 2020 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.
# ==============================================================================
"""Python API for save and loading a dataset."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
COMPRESSION_GZIP = "GZIP"
COMPRESSION_SNAPPY = "NONE"
DATASET_SPEC_FILENAME = "dataset_spec.pb"
@tf_export("data.experimental.save", v1=[])
@deprecation.deprecated(None, "Use `tf.data.Dataset.save(...)` instead.")
def save(dataset,
path,
compression=None,
shard_func=None,
checkpoint_args=None):
"""Saves the content of the given dataset.
Example usage:
>>> import tempfile
>>> path = os.path.join(tempfile.gettempdir(), "saved_data")
>>> # Save a dataset
>>> dataset = tf.data.Dataset.range(2)
>>> tf.data.experimental.save(dataset, path)
>>> new_dataset = tf.data.experimental.load(path)
>>> for elem in new_dataset:
... print(elem)
tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(1, shape=(), dtype=int64)
The saved dataset is saved in multiple file "shards". By default, the dataset
output is divided to shards in a round-robin fashion but custom sharding can
be specified via the `shard_func` function. For example, you can save the
dataset to using a single shard as follows:
```python
dataset = make_dataset()
def custom_shard_func(element):
return np.int64(0)
dataset = tf.data.experimental.save(
path="/path/to/data", ..., shard_func=custom_shard_func)
```
To enable checkpointing, pass in `checkpoint_args` to the `save` method
as follows:
```python
dataset = tf.data.Dataset.range(100)
save_dir = "..."
checkpoint_prefix = "..."
step_counter = tf.Variable(0, trainable=False)
checkpoint_args = {
"checkpoint_interval": 50,
"step_counter": step_counter,
"directory": checkpoint_prefix,
"max_to_keep": 20,
}
dataset.save(dataset, save_dir, checkpoint_args=checkpoint_args)
```
NOTE: The directory layout and file format used for saving the dataset is
considered an implementation detail and may change. For this reason, datasets
saved through `tf.data.experimental.save` should only be consumed through
`tf.data.experimental.load`, which is guaranteed to be backwards compatible.
Args:
dataset: The dataset to save.
path: Required. A directory to use for saving the dataset.
compression: Optional. The algorithm to use to compress data when writing
it. Supported options are `GZIP` and `NONE`. Defaults to `NONE`.
shard_func: Optional. A function to control the mapping of dataset elements
to file shards. The function is expected to map elements of the input
dataset to int64 shard IDs. If present, the function will be traced and
executed as graph computation.
checkpoint_args: Optional args for checkpointing which will be passed into
the `tf.train.CheckpointManager`. If `checkpoint_args` are not specified,
then checkpointing will not be performed. The `save()` implementation
creates a `tf.train.Checkpoint` object internally, so users should not
set the `checkpoint` argument in `checkpoint_args`.
Returns:
An operation which when executed performs the save. When writing
checkpoints, returns None. The return value is useful in unit tests.
Raises:
ValueError if `checkpoint` is passed into `checkpoint_args`.
"""
return dataset.save(path, compression, shard_func, checkpoint_args)
@tf_export("data.experimental.load", v1=[])
@deprecation.deprecated(None, "Use `tf.data.Dataset.load(...)` instead.")
def load(path, element_spec=None, compression=None, reader_func=None):
"""Loads a previously saved dataset.
Example usage:
>>> import tempfile
>>> path = os.path.join(tempfile.gettempdir(), "saved_data")
>>> # Save a dataset
>>> dataset = tf.data.Dataset.range(2)
>>> tf.data.experimental.save(dataset, path)
>>> new_dataset = tf.data.experimental.load(path)
>>> for elem in new_dataset:
... print(elem)
tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(1, shape=(), dtype=int64)
If the default option of sharding the saved dataset was used, the element
order of the saved dataset will be preserved when loading it.
The `reader_func` argument can be used to specify a custom order in which
elements should be loaded from the individual shards. The `reader_func` is
expected to take a single argument -- a dataset of datasets, each containing
elements of one of the shards -- and return a dataset of elements. For
example, the order of shards can be shuffled when loading them as follows:
```python
def custom_reader_func(datasets):
datasets = datasets.shuffle(NUM_SHARDS)
return datasets.interleave(lambda x: x, num_parallel_calls=AUTOTUNE)
dataset = tf.data.experimental.load(
path="/path/to/data", ..., reader_func=custom_reader_func)
```
Args:
path: Required. A path pointing to a previously saved dataset.
element_spec: Optional. A nested structure of `tf.TypeSpec` objects matching
the structure of an element of the saved dataset and specifying the type
of individual element components. If not provided, the nested structure of
`tf.TypeSpec` saved with the saved dataset is used. Note that this
argument is required in graph mode.
compression: Optional. The algorithm to use to decompress the data when
reading it. Supported options are `GZIP` and `NONE`. Defaults to `NONE`.
reader_func: Optional. A function to control how to read data from shards.
If present, the function will be traced and executed as graph computation.
Returns:
A `tf.data.Dataset` instance.
Raises:
FileNotFoundError: If `element_spec` is not specified and the saved nested
structure of `tf.TypeSpec` can not be located with the saved dataset.
ValueError: If `element_spec` is not specified and the method is executed
in graph mode.
"""
return dataset_ops.Dataset.load(path, element_spec, compression, reader_func)
tensorflow/python/data/experimental/ops/iterator_model_ops.py
+56
View File
@@ -0,0 +1,56 @@
# Copyright 2024 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.
# ==============================================================================
"""Ops for dealing with the analytical model inside of the iterator."""
from tensorflow.core.framework import model_pb2
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.eager import context
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.get_model_proto")
def get_model_proto(iterator) -> model_pb2.ModelProto:
"""Gets the analytical model inside of `iterator` as `model_pb2.ModelProto`.
Args:
iterator: An `iterator_ops.OwnedIterator` or `dataset_ops.NumpyIterator`
Returns:
The model inside of this iterator as a model proto.
Raises:
NotFoundError: If this iterator's autotune is not enabled.
"""
if isinstance(iterator, iterator_ops.OwnedIterator):
iterator_resource = iterator._iterator_resource # pylint: disable=protected-access
elif isinstance(iterator, dataset_ops.NumpyIterator):
iterator_resource = iterator._iterator._iterator_resource # pylint: disable=protected-access
else:
raise ValueError("Only supports `tf.data.Iterator`-typed `iterator`.")
if not context.executing_eagerly():
raise ValueError(
f"{get_model_proto.__name__} is not supported in graph mode."
)
model_proto_string_tensor = ged_ops.iterator_get_model_proto(
iterator_resource
)
model_proto_bytes = model_proto_string_tensor.numpy()
return model_pb2.ModelProto.FromString(model_proto_bytes)
tensorflow/python/data/experimental/ops/iterator_ops.py
+99
View File
@@ -0,0 +1,99 @@
# Copyright 2017 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.
# ==============================================================================
"""Iterator ops."""
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import options as options_lib
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
def _convert_external_state_policy_to_enum(external_state_policy):
if isinstance(external_state_policy, options_lib.ExternalStatePolicy):
return external_state_policy
if external_state_policy == "warn":
return options_lib.ExternalStatePolicy.WARN
if external_state_policy == "ignore":
return options_lib.ExternalStatePolicy.IGNORE
if external_state_policy == "fail":
return options_lib.ExternalStatePolicy.FAIL
raise ValueError(
f"Invalid `ExternalStatePolicy.` Supported values include 'warn', "
f"'ignore', and 'fail.' Received {external_state_policy}."
)
@tf_export("data.experimental.make_saveable_from_iterator")
@deprecation.deprecated(
None, "`make_saveable_from_iterator` is intended for use in TF1 with "
"`tf.compat.v1.Saver`. In TF2, use `tf.train.Checkpoint` instead.")
def make_saveable_from_iterator(iterator, external_state_policy=None):
"""Returns a SaveableObject for saving/restoring iterator state using Saver.
Args:
iterator: Iterator.
external_state_policy: A string that identifies how to handle input
pipelines that depend on external state. Possible values are
'ignore': The external state is silently ignored.
'warn': The external state is ignored, logging a warning.
'fail': The operation fails upon encountering external state.
By default we set it to 'fail'.
Returns:
A SaveableObject for saving/restoring iterator state using Saver.
Raises:
ValueError: If iterator does not support checkpointing.
ValueError: If `external_state_policy` is not one of 'warn', 'ignore' or
'fail'.
For example:
```python
with tf.Graph().as_default():
ds = tf.data.Dataset.range(10)
iterator = ds.make_initializable_iterator()
# Build the iterator SaveableObject.
saveable_obj = tf.data.experimental.make_saveable_from_iterator(iterator)
# Add the SaveableObject to the SAVEABLE_OBJECTS collection so
# it can be automatically saved using Saver.
tf.compat.v1.add_to_collection(tf.GraphKeys.SAVEABLE_OBJECTS, saveable_obj)
saver = tf.compat.v1.train.Saver()
while continue_training:
... Perform training ...
if should_save_checkpoint:
saver.save()
```
Note: When restoring the iterator, the existing iterator state is completely
discarded. This means that any changes you may have made to the Dataset
graph will be discarded as well! This includes the new Dataset graph
that you may have built during validation. So, while running validation,
make sure to run the initializer for the validation input pipeline after
restoring the checkpoint.
Note: Not all iterators support checkpointing yet. Attempting to save the
state of an unsupported iterator will throw an error.
"""
if external_state_policy is None:
external_state_policy = "fail"
policy_enum = _convert_external_state_policy_to_enum(external_state_policy)
return iterator_ops._IteratorSaveable( # pylint: disable=protected-access
iterator._iterator_resource, # pylint: disable=protected-access
iterator._iterator_resource.name, # pylint: disable=protected-access
external_state_policy=policy_enum)
tensorflow/python/data/experimental/ops/lookup_ops.py
+238
View File
@@ -0,0 +1,238 @@
# Copyright 2017 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.
#==============================================================================
"""Lookup operations."""
from tensorflow.python.data.experimental.ops.cardinality import assert_cardinality
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.util.tf_export import tf_export
def _check_table_initializer_element_spec(element_spec):
"""Raises an error if the given table initializer element spec is invalid."""
base_error = ("Datasets used to initialize lookup tables must "
"produce elements in the form (key, value), where "
"the keys and values are scalar tensors. ")
specific_error = None
if len(element_spec) != 2:
raise ValueError(base_error + "However, the given dataset produces "
f"{len(element_spec)} components instead of two "
"(key, value) components. Full dataset element spec: "
f"{element_spec}.")
if not isinstance(element_spec[0], tensor.TensorSpec):
raise ValueError(base_error + "However, the given dataset produces "
f"non-Tensor keys of type {type(element_spec[0])}.")
if not isinstance(element_spec[1], tensor.TensorSpec):
raise ValueError(base_error + "However, the given dataset produces "
f"non-Tensor values of type {type(element_spec[1])}.")
if element_spec[0].shape.rank not in (None, 0):
raise ValueError(
base_error + "However, the given dataset produces "
f"non-scalar key Tensors of rank {element_spec[0].shape.rank}.")
if element_spec[1].shape.rank not in (None, 0):
raise ValueError(
base_error + "However, the given dataset produces "
f"non-scalar value Tensors of rank {element_spec[1].shape.rank}.")
@tf_export("data.experimental.DatasetInitializer")
class DatasetInitializer(lookup_ops.TableInitializerBase):
"""Creates a table initializer from a `tf.data.Dataset`.
Sample usage:
>>> keys = tf.data.Dataset.range(100)
>>> values = tf.data.Dataset.range(100).map(
... lambda x: tf.strings.as_string(x * 2))
>>> ds = tf.data.Dataset.zip((keys, values))
>>> init = tf.data.experimental.DatasetInitializer(ds)
>>> table = tf.lookup.StaticHashTable(init, "")
>>> table.lookup(tf.constant([0, 1, 2], dtype=tf.int64)).numpy()
array([b'0', b'2', b'4'], dtype=object)
Attributes:
dataset: A `tf.data.Dataset` object that produces tuples of scalars. The
first scalar is treated as a key and the second as value.
Raises: ValueError if `dataset` doesn't conform to specifications.
"""
def __init__(self, dataset):
"""Creates a table initializer from a `tf.data.Dataset`.
Args:
dataset: A `tf.data.Dataset` object that produces tuples of scalars. The
first scalar is treated as a key and the second as value.
Raises: ValueError if `dataset` doesn't conform to specifications.
Returns: A `DatasetInitializer` object
"""
# Assert that the dataset element spec is a tuple of TensorSpecs where
# each tensor is a scalar.
self.dataset = dataset
elem_spec = self.dataset.element_spec
_check_table_initializer_element_spec(elem_spec)
key_type = elem_spec[0].dtype
value_type = elem_spec[1].dtype
super(DatasetInitializer, self).__init__(key_type, value_type)
def initialize(self, table):
lookup_ops.check_table_dtypes(table, self._key_dtype, self._value_dtype)
init_op = ged_ops.initialize_table_from_dataset(
table.resource_handle, self.dataset._variant_tensor) # pylint: disable=protected-access
ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op)
return init_op
@tf_export("data.experimental.table_from_dataset")
def table_from_dataset(dataset=None,
num_oov_buckets=0,
vocab_size=None,
default_value=None,
hasher_spec=lookup_ops.FastHashSpec,
key_dtype=dtypes.string,
name=None):
"""Returns a lookup table based on the given dataset.
This operation constructs a lookup table based on the given dataset of pairs
of (key, value).
Any lookup of an out-of-vocabulary token will return a bucket ID based on its
hash if `num_oov_buckets` is greater than zero. Otherwise it is assigned the
`default_value`.
The bucket ID range is
`[vocabulary size, vocabulary size + num_oov_buckets - 1]`.
Sample Usages:
>>> keys = tf.data.Dataset.range(100)
>>> values = tf.data.Dataset.range(100).map(
... lambda x: tf.strings.as_string(x * 2))
>>> ds = tf.data.Dataset.zip((keys, values))
>>> table = tf.data.experimental.table_from_dataset(
... ds, default_value='n/a', key_dtype=tf.int64)
>>> table.lookup(tf.constant([0, 1, 2], dtype=tf.int64)).numpy()
array([b'0', b'2', b'4'], dtype=object)
Args:
dataset: A dataset containing (key, value) pairs.
num_oov_buckets: The number of out-of-vocabulary buckets.
vocab_size: Number of the elements in the vocabulary, if known.
default_value: The value to use for out-of-vocabulary feature values.
Defaults to -1.
hasher_spec: A `HasherSpec` to specify the hash function to use for
assignation of out-of-vocabulary buckets.
key_dtype: The `key` data type.
name: A name for this op (optional).
Returns:
The lookup table based on the given dataset.
Raises:
ValueError: If
* `dataset` does not contain pairs
* The 2nd item in the `dataset` pairs has a dtype which is incompatible
with `default_value`
* `num_oov_buckets` is negative
* `vocab_size` is not greater than zero
* The `key_dtype` is not integer or string
"""
elem_spec = dataset.element_spec
_check_table_initializer_element_spec(elem_spec)
if default_value is None:
default_value = -1
if not (elem_spec[1].dtype.is_integer or elem_spec[1].dtype.is_floating):
raise ValueError("`default_value` must be specified when creating a "
"table from a dataset that produces values of type "
f"{elem_spec[1].dtype}.")
if num_oov_buckets < 0:
raise ValueError("`num_oov_buckets` must be greater than or equal to 0, "
f"got {num_oov_buckets}.")
if (not isinstance(vocab_size, tensor.Tensor) and vocab_size is not None and
vocab_size < 1):
raise ValueError(f"`vocab_size` must be greater than 0, got {vocab_size}.")
if (not key_dtype.is_integer) and (dtypes.string != key_dtype.base_dtype):
raise TypeError("`key_dtype` must be either an integer or string type, "
f"but got {key_dtype}")
if vocab_size is not None:
if isinstance(vocab_size, tensor.Tensor):
vocab_size = math_ops.cast(vocab_size, dtypes.int64)
dataset = dataset.take(vocab_size)
dataset = dataset.apply(assert_cardinality(vocab_size))
with ops.name_scope(name, "string_to_index"):
initializer = DatasetInitializer(dataset)
with ops.name_scope(None, "hash_table"):
table = lookup_ops.StaticHashTableV1(initializer, default_value)
if num_oov_buckets:
table = lookup_ops.IdTableWithHashBuckets(
table,
num_oov_buckets=num_oov_buckets,
hasher_spec=hasher_spec,
key_dtype=key_dtype)
return table
@tf_export("data.experimental.index_table_from_dataset")
def index_table_from_dataset(dataset=None,
num_oov_buckets=0,
vocab_size=None,
default_value=-1,
hasher_spec=lookup_ops.FastHashSpec,
key_dtype=dtypes.string,
name=None):
"""Returns an index lookup table based on the given dataset.
This operation constructs a lookup table based on the given dataset of keys.
Any lookup of an out-of-vocabulary token will return a bucket ID based on its
hash if `num_oov_buckets` is greater than zero. Otherwise it is assigned the
`default_value`.
The bucket ID range is
`[vocabulary size, vocabulary size + num_oov_buckets - 1]`.
Sample Usages:
>>> ds = tf.data.Dataset.range(100).map(lambda x: tf.strings.as_string(x * 2))
>>> table = tf.data.experimental.index_table_from_dataset(
... ds, key_dtype=dtypes.int64)
>>> table.lookup(tf.constant(['0', '2', '4'], dtype=tf.string)).numpy()
array([0, 1, 2])
Args:
dataset: A dataset of keys.
num_oov_buckets: The number of out-of-vocabulary buckets.
vocab_size: Number of the elements in the vocabulary, if known.
default_value: The value to use for out-of-vocabulary feature values.
Defaults to -1.
hasher_spec: A `HasherSpec` to specify the hash function to use for
assignation of out-of-vocabulary buckets.
key_dtype: The `key` data type.
name: A name for this op (optional).
Returns:
The lookup table based on the given dataset.
Raises:
ValueError: If
* `num_oov_buckets` is negative
* `vocab_size` is not greater than zero
* The `key_dtype` is not integer or string
"""
return table_from_dataset(dataset.enumerate().map(lambda v, k: (k, v)),
num_oov_buckets, vocab_size, default_value,
hasher_spec, key_dtype, name)
tensorflow/python/data/experimental/ops/map_defun.py
+65
View File
@@ -0,0 +1,65 @@
# 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 optimizing `tf.data` pipelines."""
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import gen_dataset_ops
def map_defun(fn,
elems,
output_dtypes,
output_shapes,
max_intra_op_parallelism=1):
"""Map a function on the list of tensors unpacked from `elems` on dimension 0.
Args:
fn: A function (`function.defun`) that takes a list of tensors and returns
another list of tensors. The output list has the same types as
output_dtypes. The elements of the output list have the same dimension 0
as `elems`, and the remaining dimensions correspond to those of
`fn_output_shapes`.
elems: A list of tensors.
output_dtypes: A list of dtypes corresponding to the output types of the
function.
output_shapes: A list of `TensorShape`s corresponding to the output shapes
from each invocation of the function on slices of inputs.
max_intra_op_parallelism: An integer. If positive, sets the max parallelism
limit of each function call to this.
Raises:
ValueError: if any of the inputs are malformed.
Returns:
A list of `Tensor` objects with the same types as `output_dtypes`.
"""
if not isinstance(elems, list):
raise ValueError(f"`elems` must be a list of tensors, but was {elems}.")
if not isinstance(output_dtypes, list):
raise ValueError("`output_dtypes` must be a list of `tf.DType` objects, "
f"but was {output_dtypes}.")
if not isinstance(output_shapes, list):
raise ValueError("`output_shapes` must be a list of `tf.TensorShape` "
f"objects, but was {output_shapes}.")
concrete_fn = fn.get_concrete_function() # pylint: disable=protected-access
# TODO(shivaniagrawal/rachelim): what about functions created without
# input_signature.
elems = [ops.convert_to_tensor(e) for e in elems]
output_shapes = [tensor_shape.TensorShape(s) for s in output_shapes]
return gen_dataset_ops.map_defun(elems, concrete_fn.captured_inputs,
output_dtypes, output_shapes, concrete_fn,
max_intra_op_parallelism)
tensorflow/python/data/experimental/ops/matching_files.py
+35
View File
@@ -0,0 +1,35 @@
# 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 matching input filenames."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
class MatchingFilesDataset(dataset_ops.DatasetSource):
"""A `Dataset` that list the files according to the input patterns."""
def __init__(self, patterns):
self._patterns = ops.convert_to_tensor(
patterns, dtype=dtypes.string, name="patterns")
variant_tensor = ged_ops.matching_files_dataset(self._patterns)
super(MatchingFilesDataset, self).__init__(variant_tensor)
@property
def element_spec(self):
return tensor_spec.TensorSpec([], dtypes.string)
tensorflow/python/data/experimental/ops/pad_to_cardinality.py
+105
View File
@@ -0,0 +1,105 @@
# Copyright 2023 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.
# ==============================================================================
"""The implementation of `tf.data.experimental.pad_to_cardinality`."""
from collections.abc import Mapping
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import context
from tensorflow.python.ops import array_ops
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.pad_to_cardinality")
def pad_to_cardinality(cardinality, mask_key="valid"):
"""Pads a dataset with fake elements to reach the desired cardinality.
The dataset to pad must have a known and finite cardinality and contain
dictionary elements. The `mask_key` will be added to differentiate between
real and padding elements -- real elements will have a `<mask_key>=True` entry
while padding elements will have a `<mask_key>=False` entry.
Example usage:
ds = tf.data.Dataset.from_tensor_slices({'a': [1, 2]})
ds = ds.apply(tf.data.experimental.pad_to_cardinality(3))
list(ds.as_numpy_iterator())
[{'a': 1, 'valid': True}, {'a': 2, 'valid': True}, {'a': 0, 'valid': False}]
This can be useful, e.g. during eval, when partial batches are undesirable but
it is also important not to drop any data.
```
ds = ...
# Round up to the next full batch.
target_cardinality = -(-ds.cardinality() // batch_size) * batch_size
ds = ds.apply(tf.data.experimental.pad_to_cardinality(target_cardinality))
# Set `drop_remainder` so that batch shape will be known statically. No data
# will actually be dropped since the batch size divides the cardinality.
ds = ds.batch(batch_size, drop_remainder=True)
```
Args:
cardinality: The cardinality to pad the dataset to.
mask_key: The key to use for identifying real vs padding elements.
Returns:
A dataset transformation that can be applied via `Dataset.apply()`.
"""
def make_filler_dataset(ds):
padding = cardinality - ds.cardinality()
filler_element = nest.map_structure(
lambda spec: array_ops.zeros(spec.shape, spec.dtype), ds.element_spec
)
filler_element[mask_key] = False
filler_dataset = dataset_ops.Dataset.from_tensors(filler_element)
filler_dataset = filler_dataset.repeat(padding)
return filler_dataset
def apply_valid_mask(x):
x[mask_key] = True
return x
def _apply_fn(dataset):
# The cardinality tensor is unknown during tracing, so we only check it
# in eager mode.
if context.executing_eagerly():
if dataset.cardinality() < 0:
raise ValueError(
"The dataset passed into `pad_to_cardinality` must "
"have a known cardinalty, but has cardinality "
f"{dataset.cardinality()}"
)
if dataset.cardinality() > cardinality:
raise ValueError(
"The dataset passed into `pad_to_cardinality` must "
"have a cardinalty less than the target cardinality "
f"({cardinality}), but has cardinality "
f"{dataset.cardinality()}"
)
if not isinstance(dataset.element_spec, Mapping):
raise ValueError(
"`pad_to_cardinality` requires its input dataset to "
"be a dictionary."
)
filler = make_filler_dataset(dataset)
dataset = dataset.map(apply_valid_mask)
dataset = dataset.concatenate(filler)
return dataset
return _apply_fn
tensorflow/python/data/experimental/ops/parsing_ops.py
+161
View File
@@ -0,0 +1,161 @@
# 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 `dataset` API for parsing example."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.ops import parsing_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
class _ParseExampleDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that parses `example` dataset into a `dict` dataset."""
def __init__(self, input_dataset, features, num_parallel_calls,
deterministic):
self._input_dataset = input_dataset
if not structure.are_compatible(
input_dataset.element_spec,
tensor_spec.TensorSpec([None], dtypes.string)):
raise TypeError("Input dataset should be a dataset of vectors of "
f"strings. Instead it is `{input_dataset.element_spec}`.")
self._num_parallel_calls = num_parallel_calls
if deterministic is None:
self._deterministic = "default"
elif deterministic:
self._deterministic = "true"
else:
self._deterministic = "false"
# pylint: disable=protected-access
self._features = parsing_ops._prepend_none_dimension(features)
params = parsing_ops._ParseOpParams.from_features(self._features, [
parsing_ops.VarLenFeature, parsing_ops.SparseFeature,
parsing_ops.FixedLenFeature, parsing_ops.FixedLenSequenceFeature,
parsing_ops.RaggedFeature
])
# pylint: enable=protected-access
self._sparse_keys = params.sparse_keys
self._sparse_types = params.sparse_types
self._ragged_keys = params.ragged_keys
self._ragged_value_types = params.ragged_value_types
self._ragged_split_types = params.ragged_split_types
self._dense_keys = params.dense_keys
self._dense_defaults = params.dense_defaults_vec
self._dense_shapes = params.dense_shapes_as_proto
self._dense_types = params.dense_types
input_dataset_shape = dataset_ops.get_legacy_output_shapes(
self._input_dataset)
self._element_spec = {}
for (key, value_type) in zip(params.sparse_keys, params.sparse_types):
self._element_spec[key] = sparse_tensor.SparseTensorSpec(
input_dataset_shape.concatenate([None]), value_type)
for (key, value_type, dense_shape) in zip(params.dense_keys,
params.dense_types,
params.dense_shapes):
self._element_spec[key] = tensor_spec.TensorSpec(
input_dataset_shape.concatenate(dense_shape), value_type)
for (key, value_type, splits_type) in zip(params.ragged_keys,
params.ragged_value_types,
params.ragged_split_types):
self._element_spec[key] = ragged_tensor.RaggedTensorSpec(
input_dataset_shape.concatenate([None]), value_type, 1, splits_type)
variant_tensor = (
gen_experimental_dataset_ops.parse_example_dataset_v2(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._num_parallel_calls,
self._dense_defaults,
self._sparse_keys,
self._dense_keys,
self._sparse_types,
self._dense_shapes,
deterministic=self._deterministic,
ragged_keys=self._ragged_keys,
ragged_value_types=self._ragged_value_types,
ragged_split_types=self._ragged_split_types,
**self._flat_structure))
super(_ParseExampleDataset, self).__init__(input_dataset, variant_tensor)
@property
def element_spec(self):
return self._element_spec
@tf_export("data.experimental.parse_example_dataset")
@deprecation.deprecated(
None, "Use `tf.data.Dataset.map(tf.io.parse_example(...))` instead.")
def parse_example_dataset(features, num_parallel_calls=1, deterministic=None):
"""A transformation that parses `Example` protos into a `dict` of tensors.
Parses a number of serialized `Example` protos given in `serialized`. We refer
to `serialized` as a batch with `batch_size` many entries of individual
`Example` protos.
This op parses serialized examples into a dictionary mapping keys to `Tensor`,
`SparseTensor`, and `RaggedTensor` objects. `features` is a dict from keys to
`VarLenFeature`, `RaggedFeature`, `SparseFeature`, and `FixedLenFeature`
objects. Each `VarLenFeature` and `SparseFeature` is mapped to a
`SparseTensor`; each `RaggedFeature` is mapped to a `RaggedTensor`; and each
`FixedLenFeature` is mapped to a `Tensor`. See `tf.io.parse_example` for more
details about feature dictionaries.
Args:
features: A `dict` mapping feature keys to `FixedLenFeature`,
`VarLenFeature`, `RaggedFeature`, and `SparseFeature` values.
num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
representing the number of parsing processes to call in parallel.
deterministic: (Optional.) A boolean controlling whether determinism
should be traded for performance by allowing elements to be produced out
of order if some parsing calls complete faster than others. If
`deterministic` is `None`, the
`tf.data.Options.deterministic` dataset option (`True` by default) is used
to decide whether to produce elements deterministically.
Returns:
A dataset transformation function, which can be passed to
`tf.data.Dataset.apply`.
Raises:
ValueError: if features argument is None.
"""
if features is None:
raise ValueError("Argument `features` is required, but not specified.")
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
out_dataset = _ParseExampleDataset(dataset, features, num_parallel_calls,
deterministic)
if any(
isinstance(feature, parsing_ops.SparseFeature) or
isinstance(feature, parsing_ops.RaggedFeature)
for feature in features.values()):
# pylint: disable=protected-access
# pylint: disable=g-long-lambda
out_dataset = out_dataset.map(
lambda x: parsing_ops._construct_tensors_for_composite_features(
features, x),
num_parallel_calls=num_parallel_calls)
return out_dataset
return _apply_fn
tensorflow/python/data/experimental/ops/prefetching_ops.py
+287
View File
@@ -0,0 +1,287 @@
# Copyright 2017 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.
# ==============================================================================
"""Python wrapper for prefetching_ops."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import structured_function
from tensorflow.python.data.util import structure
from tensorflow.python.eager import def_function
from tensorflow.python.framework import device as framework_device
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import gen_dataset_ops
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.prefetch_to_device")
def prefetch_to_device(device, buffer_size=None):
"""A transformation that prefetches dataset values to the given `device`.
NOTE: Although the transformation creates a `tf.data.Dataset`, the
transformation must be the final `Dataset` in the input pipeline.
For example,
>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
>>> dataset = dataset.apply(tf.data.experimental.prefetch_to_device("/cpu:0"))
>>> for element in dataset:
... print(f'Tensor {element} is on device {element.device}')
Tensor 1 is on device /job:localhost/replica:0/task:0/device:CPU:0
Tensor 2 is on device /job:localhost/replica:0/task:0/device:CPU:0
Tensor 3 is on device /job:localhost/replica:0/task:0/device:CPU:0
Args:
device: A string. The name of a device to which elements will be prefetched.
buffer_size: (Optional.) The number of elements to buffer on `device`.
Defaults to an automatically chosen value.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.apply(
copy_to_device(target_device=device)).prefetch(buffer_size)
return _apply_fn
@tf_export("data.experimental.copy_to_device")
def copy_to_device(target_device, source_device="/cpu:0"):
"""A transformation that copies dataset elements to the given `target_device`.
Args:
target_device: The name of a device to which elements will be copied.
source_device: The original device on which `input_dataset` will be placed.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return _CopyToDeviceDataset(
dataset, target_device=target_device, source_device=source_device)
return _apply_fn
# TODO(rohanj): Use the _input_hostmem attr on the RemoteCall ops to indicate
# all inputs to the Op are in host memory, thereby avoiding some unnecessary
# Sends and Recvs.
class _CopyToDeviceDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that copies elements to another device."""
def __init__(self, input_dataset, target_device, source_device="/cpu:0"):
"""Constructs a _CopyToDeviceDataset.
Args:
input_dataset: `Dataset` to be copied
target_device: The name of the device to which elements would be copied.
source_device: Device where input_dataset would be placed.
"""
self._input_dataset = input_dataset._apply_debug_options() # pylint: disable=protected-access
self._target_device = target_device
spec = framework_device.DeviceSpec().from_string(self._target_device)
self._is_gpu_target = (spec.device_type == "GPU")
self._source_device_string = source_device
self._source_device = ops.convert_to_tensor(source_device)
wrap_ds_variant = gen_dataset_ops.wrap_dataset_variant(
self._input_dataset._variant_tensor) # pylint: disable=protected-access
@def_function.function()
def _init_func():
"""Creates an iterator for the input dataset.
Returns:
A `string` tensor that encapsulates the iterator created.
"""
ds_variant = gen_dataset_ops.unwrap_dataset_variant(wrap_ds_variant)
resource = gen_dataset_ops.anonymous_iterator(
**self._input_dataset._flat_structure) # pylint: disable=protected-access
with ops.control_dependencies(
[gen_dataset_ops.make_iterator(ds_variant, resource)]):
return gen_dataset_ops.iterator_to_string_handle(resource)
init_func_concrete = _init_func.get_concrete_function() # pylint: disable=protected-access
@def_function.function()
def _remote_init_func():
return functional_ops.remote_call(
target=self._source_device,
args=init_func_concrete.captured_inputs,
Tout=[dtypes.string],
f=init_func_concrete)
self._init_func = _remote_init_func.get_concrete_function() # pylint: disable=protected-access
self._init_captured_args = self._init_func.captured_inputs
@def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _next_func(string_handle):
"""Calls get_next for created iterator.
Args:
string_handle: An iterator string handle created by _init_func
Returns:
The elements generated from `input_dataset`
"""
with ops.device(self._source_device_string):
iterator = iterator_ops.Iterator.from_string_handle(
string_handle,
dataset_ops.get_legacy_output_types(self),
dataset_ops.get_legacy_output_shapes(self),
dataset_ops.get_legacy_output_classes(self))
return structure.to_tensor_list(self.element_spec, iterator.get_next())
next_func_concrete = _next_func.get_concrete_function() # pylint: disable=protected-access
@def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.string)],
experimental_attributes={"experimental_ints_on_device": True})
def _remote_next_func(string_handle):
return functional_ops.remote_call(
target=self._source_device,
args=[string_handle] + next_func_concrete.captured_inputs,
Tout=self._input_dataset._flat_types, # pylint: disable=protected-access
f=next_func_concrete)
self._next_func = _remote_next_func.get_concrete_function()
self._next_captured_args = self._next_func.captured_inputs
@def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _finalize_func(string_handle):
"""Destroys the iterator resource created.
Args:
string_handle: An iterator string handle created by _init_func
Returns:
Tensor constant 0
"""
iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2(
string_handle,
**self._input_dataset._flat_structure) # pylint: disable=protected-access
with ops.control_dependencies([
resource_variable_ops.destroy_resource_op(
iterator_resource, ignore_lookup_error=True)]):
return array_ops.constant(0, dtypes.int64)
finalize_func_concrete = _finalize_func.get_concrete_function() # pylint: disable=protected-access
@def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _remote_finalize_func(string_handle):
return functional_ops.remote_call(
target=self._source_device,
args=[string_handle] + finalize_func_concrete.captured_inputs,
Tout=[dtypes.int64],
f=finalize_func_concrete)
self._finalize_func = _remote_finalize_func.get_concrete_function( # pylint: disable=protected-access
)
self._finalize_captured_args = self._finalize_func.captured_inputs
g = ops.get_default_graph()
self._init_func.add_to_graph(g)
self._next_func.add_to_graph(g)
self._finalize_func.add_to_graph(g)
# pylint: enable=protected-scope
with ops.device(self._target_device):
variant_tensor = gen_dataset_ops.generator_dataset(
self._init_captured_args,
self._next_captured_args,
self._finalize_captured_args,
init_func=self._init_func,
next_func=self._next_func,
finalize_func=self._finalize_func,
**self._input_dataset._flat_structure) # pylint: disable=protected-access
super(_CopyToDeviceDataset, self).__init__(input_dataset, variant_tensor)
# The one_shot_iterator implementation needs a 0 arg _make_dataset function
# that thereby captures all the inputs required to create the dataset. Since
# there are strings that are inputs to the GeneratorDataset which can't be
# placed on a GPU, this fails for the GPU case. Therefore, disabling it for
# GPU
def make_one_shot_iterator(self):
if self._is_gpu_target:
raise ValueError(
"`make_one_shot_iterator` is not compatible with GPU execution. "
"Please use `Dataset.make_initializable_iterator()` instead."
)
else:
return super(_CopyToDeviceDataset, self).make_one_shot_iterator()
class _MapOnGpuDataset(dataset_ops.UnaryDataset):
"""A `Dataset` that maps a function over elements in its using a GPU."""
def __init__(self, input_dataset, map_func, use_inter_op_parallelism=True):
"""See `Dataset.map()` for details."""
self._input_dataset = input_dataset
self._use_inter_op_parallelism = use_inter_op_parallelism
self._map_func = structured_function.StructuredFunctionWrapper(
map_func,
self._transformation_name(),
dataset=input_dataset,
defun_kwargs={"experimental_ints_on_device": True})
variant_tensor = ged_ops.experimental_map_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._map_func.function.captured_inputs,
f=self._map_func.function,
use_inter_op_parallelism=self._use_inter_op_parallelism,
**self._flat_structure)
super(_MapOnGpuDataset, self).__init__(input_dataset, variant_tensor)
def _functions(self):
return [self._map_func]
@property
def element_spec(self):
return self._map_func.output_structure
def _transformation_name(self):
return "map_on_gpu()"
def map_on_gpu(map_func):
"""Maps `map_func` across the elements of this dataset.
NOTE: This is a highly experimental version of `tf.data.Dataset.map` that runs
`map_func` on GPU. It must be used after applying the
`tf.data.experimental.copy_to_device` transformation with a GPU device
argument.
Args:
map_func: A function mapping a nested structure of tensors (having shapes
and types defined by `self.output_shapes` and `self.output_types`) to
another nested structure of tensors.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return _MapOnGpuDataset(dataset, map_func)
return _apply_fn
tensorflow/python/data/experimental/ops/random_access.py
+73
View File
@@ -0,0 +1,73 @@
# Copyright 2021 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.
# ==============================================================================
"""Python API for random indexing into a dataset."""
from tensorflow.python.data.util import structure
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.at", v1=[])
def at(dataset, index):
"""Returns the element at a specific index in a datasest.
Currently, random access is supported for the following tf.data operations:
- `tf.data.Dataset.from_tensor_slices`,
- `tf.data.Dataset.from_tensors`,
- `tf.data.Dataset.shuffle`,
- `tf.data.Dataset.batch`,
- `tf.data.Dataset.shard`,
- `tf.data.Dataset.map`,
- `tf.data.Dataset.range`,
- `tf.data.Dataset.zip`,
- `tf.data.Dataset.skip`,
- `tf.data.Dataset.repeat`,
- `tf.data.Dataset.list_files`,
- `tf.data.Dataset.SSTableDataset`,
- `tf.data.Dataset.concatenate`,
- `tf.data.Dataset.enumerate`,
- `tf.data.Dataset.parallel_map`,
- `tf.data.Dataset.prefetch`,
- `tf.data.Dataset.take`,
- `tf.data.Dataset.cache` (in-memory only)
Users can use the cache operation to enable random access for any dataset,
even one comprised of transformations which are not on this list.
E.g., to get the third element of a TFDS dataset:
```python
ds = tfds.load("mnist", split="train").cache()
elem = tf.data.Dataset.experimental.at(ds, 3)
```
Args:
dataset: A `tf.data.Dataset` to determine whether it supports random access.
index: The index at which to fetch the element.
Returns:
A (nested) structure of values matching `tf.data.Dataset.element_spec`.
Raises:
UnimplementedError: If random access is not yet supported for a dataset.
"""
# pylint: disable=protected-access
return structure.from_tensor_list(
dataset.element_spec,
gen_experimental_dataset_ops.get_element_at_index(
dataset._variant_tensor,
index,
output_types=structure.get_flat_tensor_types(dataset.element_spec),
output_shapes=structure.get_flat_tensor_shapes(dataset.element_spec)))
tensorflow/python/data/experimental/ops/random_ops.py
+58
View File
@@ -0,0 +1,58 @@
# Copyright 2017 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.
# ==============================================================================
"""Datasets for random number generators."""
import functools
from tensorflow.python import tf2
from tensorflow.python.compat import v2_compat
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import random_op
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
# TODO(b/260143413): Migrate users to `tf.data.Dataset.random`.
@deprecation.deprecated(None, "Use `tf.data.Dataset.random(...)`.")
@tf_export("data.experimental.RandomDataset", v1=[])
class RandomDatasetV2(random_op._RandomDataset): # pylint: disable=protected-access
"""A `Dataset` of pseudorandom values."""
@deprecation.deprecated(None, "Use `tf.data.Dataset.random(...)`.")
@tf_export(v1=["data.experimental.RandomDataset"])
class RandomDatasetV1(dataset_ops.DatasetV1Adapter):
"""A `Dataset` of pseudorandom values."""
@functools.wraps(RandomDatasetV2.__init__)
def __init__(self, seed=None):
wrapped = RandomDatasetV2(seed)
super(RandomDatasetV1, self).__init__(wrapped)
if tf2.enabled():
RandomDataset = RandomDatasetV2
else:
RandomDataset = RandomDatasetV1
def _tf2_callback():
global RandomDataset
if tf2.enabled():
RandomDataset = RandomDatasetV2
else:
RandomDataset = RandomDatasetV1
v2_compat.register_data_v2_callback(_tf2_callback)
tensorflow/python/data/experimental/ops/readers.py
+1239
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/data/experimental/ops/resampling.py
+50
View File
@@ -0,0 +1,50 @@
# Copyright 2017 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.
# ==============================================================================
"""Resampling dataset transformations."""
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@deprecation.deprecated(None, "Use `tf.data.Dataset.rejection_resample(...)`.")
@tf_export("data.experimental.rejection_resample")
def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
"""A transformation that resamples a dataset to achieve a target distribution.
**NOTE** Resampling is performed via rejection sampling; some fraction
of the input values will be dropped.
Args:
class_func: A function mapping an element of the input dataset to a scalar
`tf.int32` tensor. Values should be in `[0, num_classes)`.
target_dist: A floating point type tensor, shaped `[num_classes]`.
initial_dist: (Optional.) A floating point type tensor, shaped
`[num_classes]`. If not provided, the true class distribution is
estimated live in a streaming fashion.
seed: (Optional.) Python integer seed for the resampler.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
return dataset.rejection_resample(
class_func=class_func,
target_dist=target_dist,
initial_dist=initial_dist,
seed=seed)
return _apply_fn
tensorflow/python/data/experimental/ops/scan_ops.py
+45
View File
@@ -0,0 +1,45 @@
# Copyright 2017 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.
# ==============================================================================
"""Scan dataset transformation."""
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@deprecation.deprecated(None, "Use `tf.data.Dataset.scan(...) instead")
@tf_export("data.experimental.scan")
def scan(initial_state, scan_func):
"""A transformation that scans a function across an input dataset.
This transformation is a stateful relative of `tf.data.Dataset.map`.
In addition to mapping `scan_func` across the elements of the input dataset,
`scan()` accumulates one or more state tensors, whose initial values are
`initial_state`.
Args:
initial_state: A nested structure of tensors, representing the initial state
of the accumulator.
scan_func: A function that maps `(old_state, input_element)` to
`(new_state, output_element)`. It must take two arguments and return a
pair of nested structures of tensors. The `new_state` must match the
structure of `initial_state`.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.scan(initial_state=initial_state, scan_func=scan_func)
return _apply_fn
tensorflow/python/data/experimental/ops/shuffle_ops.py
+272
View File
@@ -0,0 +1,272 @@
# Copyright 2017 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 shuffle ops."""
import functools
import numpy as np
from tensorflow.python.data.experimental.ops import random_access
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import random_seed
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_dataset_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
class _ShuffleAndRepeatDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that fuses `shuffle` and `repeat`."""
def __init__(self, input_dataset, buffer_size, count=None, seed=None):
self._input_dataset = input_dataset
self._buffer_size = ops.convert_to_tensor(
buffer_size, dtype=dtypes.int64, name="buffer_size")
if count is None:
self._count = constant_op.constant(-1, dtype=dtypes.int64, name="count")
else:
self._count = ops.convert_to_tensor(
count, dtype=dtypes.int64, name="count")
self._seed, self._seed2 = random_seed.get_seed(seed)
variant_tensor = gen_dataset_ops.shuffle_and_repeat_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
buffer_size=self._buffer_size,
count=self._count,
seed=self._seed,
seed2=self._seed2,
**self._flat_structure)
super(_ShuffleAndRepeatDataset, self).__init__(input_dataset,
variant_tensor)
@deprecation.deprecated(
None, "Use `tf.data.Dataset.shuffle(buffer_size, seed)` followed by "
"`tf.data.Dataset.repeat(count)`. Static tf.data optimizations will take "
"care of using the fused implementation.")
@tf_export("data.experimental.shuffle_and_repeat")
def shuffle_and_repeat(buffer_size, count=None, seed=None):
"""Shuffles and repeats a Dataset, reshuffling with each repetition.
>>> d = tf.data.Dataset.from_tensor_slices([1, 2, 3])
>>> d = d.apply(tf.data.experimental.shuffle_and_repeat(2, count=2))
>>> [elem.numpy() for elem in d] # doctest: +SKIP
[2, 3, 1, 1, 3, 2]
```python
dataset.apply(
tf.data.experimental.shuffle_and_repeat(buffer_size, count, seed))
```
produces the same output as
```python
dataset.shuffle(
buffer_size, seed=seed, reshuffle_each_iteration=True).repeat(count)
```
In each repetition, this dataset fills a buffer with `buffer_size` elements,
then randomly samples elements from this buffer, replacing the selected
elements with new elements. For perfect shuffling, set the buffer size equal
to the full size of the dataset.
For instance, if your dataset contains 10,000 elements but `buffer_size` is
set to 1,000, then `shuffle` will initially select a random element from
only the first 1,000 elements in the buffer. Once an element is selected,
its space in the buffer is replaced by the next (i.e. 1,001-st) element,
maintaining the 1,000 element buffer.
Args:
buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the maximum
number elements that will be buffered when prefetching.
count: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the number
of times the dataset should be repeated. The default behavior (if `count`
is `None` or `-1`) is for the dataset be repeated indefinitely.
seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the random
seed that will be used to create the distribution. See
`tf.random.set_seed` for behavior.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset): # pylint: disable=missing-docstring
return _ShuffleAndRepeatDataset(dataset, buffer_size, count, seed)
return _apply_fn
def _process_file_infos(file_infos):
"""Computes aggregate information about files to read.
The method collects information about the files to read, the total number of
elements, and arrays that can be used to account for elements to be skipped,
which can be specified via the "skip" and "take" keys.
To account for elements to skip, the range of each file can be divided into
three regions:
- S (elements to skip)
- T (elements to read)
- R (remainder of elements that will also be skipped)
The `thresholds` and `offsets` arrays are initialized as follows:
`thresholds = [0, T_1, T_1 + T_2, ...]` and
`offsets = [S_1, S_1 + R_1 + S_2, S_1 + R_1 + S_2 + R_2 + S_3, ...]`
This makes it possible to map an index from a contiguous range
`(0...num_elements_to_read)` to an index in the range of all elements,
skipping over elements as per the "skip" and "take" keys values. In
particular, for a given input index `X`, we find the greatest `thresholds`
value that is smaller or equal to `X`. Let `t(X)` denotes such index in the
`thresholds` array. The output index is computed as `X + offsets[t(X)]`.
Args:
file_infos: See `file_infos` argument of `index_shuffle` for details.
Returns:
A dictionary containing the following keys:
- `files`, the vector of pathnames of files to read
- `num_elements`, an integer identifying the total number of elements
- `offsets`, the vector of offsets to use for index adjustment (in case
any elements should be skipped)
- `thresholds`, the vector of thresholds to use for index adjustment (in
case any elements should be skipped)
"""
files = []
num_elements = 0
offsets = np.int64([])
offset_sum = 0
thresholds = np.int64([])
threshold_sum = 0
adjustment_needed = False
for file_info in file_infos:
files.append(file_info["path"])
skip = 0
if "skip" in file_info:
if file_info["skip"] < -1:
raise ValueError("`skip` should be greater than `-1` but got {}".format(
file_info["skip"]))
if file_info["skip"] == -1:
skip = file_info["num_elements"]
else:
skip = min(file_info["skip"], file_info["num_elements"])
take = file_info["num_elements"] - skip
if "take" in file_info:
if file_info["take"] < -1:
raise ValueError("`take` should be greater than `-1` but got {}".format(
file_info["take"]))
# `file_info["take"] == -1` is a no-op
if file_info["take"] != -1:
take = min(file_info["take"], take)
remainder = file_info["num_elements"] - skip - take
if take != file_info["num_elements"]:
adjustment_needed = True
num_elements += take
offsets = np.append(offsets, offset_sum + skip)
offset_sum += skip + remainder
thresholds = np.append(thresholds, threshold_sum)
threshold_sum += take
result = {"files": files, "num_elements": num_elements}
if adjustment_needed:
result["offsets"] = offsets
result["thresholds"] = thresholds
return result
def _adjust_index(index, thresholds, offsets):
"""Adjusts index to account for elements to be skipped."""
t_index = array_ops.shape(
array_ops.boolean_mask(
thresholds,
math_ops.less_equal(thresholds, index)))[0] - 1
return index + array_ops.gather(offsets, t_index)
# TODO(jsimsa): Expose this method in the public API. When we do, consider
# defining `FileInfo` as a public API to encapsulate the information provided
# through the `file_infos` argument.
def index_shuffle(file_infos,
reader_factory,
seed=None,
reshuffle_each_iteration=False,
num_parallel_calls=dataset_ops.AUTOTUNE):
"""Creates a (globally) shuffled dataset from the given set of files.
Unlike `tf.data.Dataset.shuffle()`, which uses an in-memory buffer to shuffle
elements of input dataset in a streaming fashion,
`tf.data.experimental.index_shuffle()` performs a global shuffle of element
indices and then reads the data in a shuffled order. The advantage of
`index_shuffle()` is that it can perform global shuffle of datasets that do
not fit into memory (as long as the array of their indices does) and that the
shuffling logic it provides is compatible with symbolic checkpointing. The
disadvantage of `index_shuffle()` is that reading data in a shuffled random
order will in general not be as efficient as reading data sequentially.
Args:
file_infos: A list of dictionaries that describe each file of the input
dataset. Each dictionary is expected to contain the "path" key, which
identifies the path of the file and the "num_elements" key, which
identifies the number of elements in the file. In addition, the "skip"
and "take" keys can be used to identify the number of elements to skip
and take respectively. By default, no elements are skipped and all
elements are taken.
reader_factory: A function that maps a sequence of filenames to an instance
of `tf.data.Dataset` that reads data from the files.
seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the random
seed that will be used to shuffle the order of elements. Default to
non-deterministic seed.
reshuffle_each_iteration: (Optional.) A `tf.bool` scalar `tf.Tensor`, that
determines whether to change the shuffle order each iteration. Defaults to
`False`.
num_parallel_calls: (Optional.) A `tf.int64` scalar `tf.Tensor`, that
determines the maximum number of random access operations to perform
in parallel. By default, the tf.data runtime uses autotuning to determine
the value dynamically.
Returns:
A `tf.data.Dataset` object, representing a globally shuffled dataset of
the input data.
"""
result = _process_file_infos(file_infos)
def sequential_index_shuffle(seeds):
dataset = dataset_ops.Dataset.range(result["num_elements"])
def read_element(dataset, index):
# 1) Shuffle the index.
shuffled_index = stateless_random_ops.index_shuffle(
index, seeds, result["num_elements"] - 1)
# 2) If needed, adjust the index to the non-contiguous range.
if "thresholds" in result and "offsets" in result:
shuffled_index = _adjust_index(shuffled_index, result["thresholds"],
result["offsets"])
# 3) Perform the read.
return random_access.at(dataset, shuffled_index)
# We evaluate `reader_factory()` eagerly to prevent the dataset from being
# created on every lookup.
map_func = functools.partial(read_element, reader_factory(result["files"]))
return dataset.map(map_func, num_parallel_calls=num_parallel_calls)
rng_ds = dataset_ops.Dataset.random(
seed=seed,
rerandomize_each_iteration=reshuffle_each_iteration)
rng_ds = rng_ds.take(2).batch(2, drop_remainder=True)
return rng_ds.flat_map(sequential_index_shuffle)
tensorflow/python/data/experimental/ops/sleep.py
+48
View File
@@ -0,0 +1,48 @@
# 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 manually injecting delays into `tf.data` pipelines."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.ops import gen_experimental_dataset_ops
class _SleepDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that sleeps before producing each upstream element."""
def __init__(self, input_dataset, sleep_microseconds):
self._input_dataset = input_dataset
self._sleep_microseconds = sleep_microseconds
variant_tensor = gen_experimental_dataset_ops.sleep_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._sleep_microseconds,
**self._flat_structure)
super(_SleepDataset, self).__init__(input_dataset, variant_tensor)
def sleep(sleep_microseconds):
"""Sleeps for `sleep_microseconds` before producing each input element.
Args:
sleep_microseconds: The number of microseconds to sleep before producing an
input element.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return _SleepDataset(dataset, sleep_microseconds)
return _apply_fn
tensorflow/python/data/experimental/ops/snapshot.py
+276
View File
@@ -0,0 +1,276 @@
# Copyright 2019 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.
# ==============================================================================
"""Dataset snapshot and related functionality."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
COMPRESSION_GZIP = "GZIP"
COMPRESSION_SNAPPY = "SNAPPY"
COMPRESSION_NONE = None
class _LegacySnapshotDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A Dataset that captures a snapshot or reads from a snapshot."""
def __init__(self,
input_dataset,
path,
compression=None,
reader_path_prefix=None,
writer_path_prefix=None,
shard_size_bytes=None,
pending_snapshot_expiry_seconds=None,
num_reader_threads=None,
reader_buffer_size=None,
num_writer_threads=None,
writer_buffer_size=None,
shuffle_on_read=None,
shuffle_seed=None,
mode=None,
snapshot_name=None):
self._compression = compression if compression is not None else ""
self._reader_path_prefix = (
reader_path_prefix if reader_path_prefix is not None else "")
self._writer_path_prefix = (
writer_path_prefix if writer_path_prefix is not None else "")
self._shard_size_bytes = (
shard_size_bytes if shard_size_bytes is not None else -1)
self._pending_snapshot_expiry_seconds = (
pending_snapshot_expiry_seconds
if pending_snapshot_expiry_seconds is not None else -1)
self._num_reader_threads = (
num_reader_threads if num_reader_threads is not None else -1)
self._reader_buffer_size = (
reader_buffer_size if reader_buffer_size is not None else -1)
self._num_writer_threads = (
num_writer_threads if num_writer_threads is not None else -1)
self._writer_buffer_size = (
writer_buffer_size if writer_buffer_size is not None else -1)
self._shuffle_on_read = (
shuffle_on_read if shuffle_on_read is not None else False)
self._mode = (mode if mode is not None else "auto")
self._snapshot_name = (snapshot_name if snapshot_name is not None else "")
self._seed, self._seed2 = random_seed.get_seed(shuffle_seed)
self._input_dataset = input_dataset
self._path = ops.convert_to_tensor(path, dtype=dtypes.string, name="path")
variant_tensor = ged_ops.snapshot_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
path=self._path,
compression=self._compression,
reader_path_prefix=self._reader_path_prefix,
writer_path_prefix=self._writer_path_prefix,
shard_size_bytes=self._shard_size_bytes,
pending_snapshot_expiry_seconds=self._pending_snapshot_expiry_seconds,
num_reader_threads=self._num_reader_threads,
reader_buffer_size=self._reader_buffer_size,
num_writer_threads=self._num_writer_threads,
writer_buffer_size=self._writer_buffer_size,
shuffle_on_read=self._shuffle_on_read,
seed=self._seed,
seed2=self._seed2,
mode=self._mode,
snapshot_name=self._snapshot_name,
**self._flat_structure)
super(_LegacySnapshotDataset, self).__init__(input_dataset, variant_tensor)
@deprecation.deprecated(None, "Use `tf.data.Dataset.snapshot(...)` instead.")
def legacy_snapshot(path,
compression=None,
reader_path_prefix=None,
writer_path_prefix=None,
shard_size_bytes=None,
pending_snapshot_expiry_seconds=None,
num_reader_threads=None,
reader_buffer_size=None,
num_writer_threads=None,
writer_buffer_size=None,
shuffle_on_read=None,
shuffle_seed=None,
mode=None,
snapshot_name=None):
"""Writes to/reads from a snapshot of a dataset.
This function attempts to determine whether a valid snapshot exists at the
`path`, and reads from the snapshot if so. If not, it will run the
preprocessing pipeline as usual, and write out a snapshot of the data
processed for future use.
Args:
path: A directory where we want to save our snapshots and/or read from a
previously saved snapshot.
compression: The type of compression to apply to the Dataset. Currently
supports "GZIP" or None. Defaults to None (no compression).
reader_path_prefix: A prefix to add to the path when reading from snapshots.
Defaults to None.
writer_path_prefix: A prefix to add to the path when writing to snapshots.
Defaults to None.
shard_size_bytes: The size of each shard to be written by the snapshot
dataset op. Defaults to 10 GiB.
pending_snapshot_expiry_seconds: How long to wait (in seconds) before the
snapshot op considers a previously unfinished snapshot to be stale.
num_reader_threads: Number of threads to parallelize reading from snapshot.
Especially useful if compression is turned on since the decompression
operation tends to be intensive. Defaults to 1. If > 1, then this might
introduce non-determinism i.e. the order in which the elements are read
from the snapshot are different from the order they're written.
reader_buffer_size: Maximum number of elements we can prefetch reading from
the snapshot. Defaults to 1. Increasing this might improve performance but
will increase memory consumption.
num_writer_threads: Number of threads to parallelize writing from snapshot.
We'll open up `num_writer_threads` files and write to them in parallel.
Especially useful if compression is turned on since the compression
operation tends to be intensive. Defaults to 1. If > 1, then this might
introduce non-determinism i.e. the order in which the elements are read
from the upstream iterator are different from the order they're written.
writer_buffer_size: Maximum number of pipeline elements to fill up the
buffer before writing them out using `num_writer_threads`.
shuffle_on_read: If this is True, then the order in which examples are
produced when reading from a snapshot will be random. Defaults to False.
shuffle_seed: Optional. If shuffle_seed is set, the random number generator
used for shuffling (when shuffle_on_read is turned on) is seeded by the
given seed. Otherwise, it is seeded by a random seed that differs for
every run.
mode: The mode at which snapshot should operate. Valid options are "auto",
"read", "write", and "passthrough". The default mode is "auto", where the
snapshot op will automatically determine what mode to operate in.
snapshot_name: If set, use the supplied string as a named snapshot name
instead of introspecting the data pipeline and automatically generating a
unique identifier for the snapshot.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return _LegacySnapshotDataset(
input_dataset=dataset,
path=path,
compression=compression,
reader_path_prefix=reader_path_prefix,
writer_path_prefix=writer_path_prefix,
shard_size_bytes=shard_size_bytes,
pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds,
num_reader_threads=num_reader_threads,
reader_buffer_size=reader_buffer_size,
num_writer_threads=num_writer_threads,
writer_buffer_size=writer_buffer_size,
shuffle_on_read=shuffle_on_read,
shuffle_seed=shuffle_seed,
mode=mode,
snapshot_name=snapshot_name)
return _apply_fn
@deprecation.deprecated(None, "Use `tf.data.Dataset.snapshot(...)`.")
@tf_export("data.experimental.snapshot")
def snapshot(path, compression="AUTO", reader_func=None, shard_func=None):
"""API to persist the output of the input dataset.
The snapshot API allows users to transparently persist the output of their
preprocessing pipeline to disk, and materialize the pre-processed data on a
different training run.
This API enables repeated preprocessing steps to be consolidated, and allows
re-use of already processed data, trading off disk storage and network
bandwidth for freeing up more valuable CPU resources and accelerator compute
time.
https://github.com/tensorflow/community/blob/master/rfcs/20200107-tf-data-snapshot.md
has detailed design documentation of this feature.
Users can specify various options to control the behavior of snapshot,
including how snapshots are read from and written to by passing in
user-defined functions to the `reader_func` and `shard_func` parameters.
`shard_func` is a user specified function that maps input elements to snapshot
shards.
Users may want to specify this function to control how snapshot files should
be written to disk. Below is an example of how a potential shard_func could
be written.
```python
dataset = ...
dataset = dataset.enumerate()
dataset = dataset.apply(tf.data.Dataset.snapshot("/path/to/snapshot/dir",
shard_func=lambda x, y: x % NUM_SHARDS, ...))
dataset = dataset.map(lambda x, y: y)
```
`reader_func` is a user specified function that accepts a single argument:
(1) a Dataset of Datasets, each representing a "split" of elements of the
original dataset. The cardinality of the input dataset matches the
number of the shards specified in the `shard_func` (see above). The function
should return a Dataset of elements of the original dataset.
Users may want specify this function to control how snapshot files should be
read from disk, including the amount of shuffling and parallelism.
Here is an example of a standard reader function a user can define. This
function enables both dataset shuffling and parallel reading of datasets:
```python
def user_reader_func(datasets):
# shuffle the datasets splits
datasets = datasets.shuffle(NUM_CORES)
# read datasets in parallel and interleave their elements
return datasets.interleave(lambda x: x, num_parallel_calls=AUTOTUNE)
dataset = dataset.apply(tf.data.Dataset.snapshot("/path/to/snapshot/dir",
reader_func=user_reader_func))
```
By default, snapshot parallelizes reads by the number of cores available on
the system, but will not attempt to shuffle the data.
Args:
path: Required. A directory to use for storing / loading the snapshot to /
from.
compression: Optional. The type of compression to apply to the snapshot
written to disk. Supported options are `GZIP`, `SNAPPY`, `AUTO` or None.
Defaults to AUTO, which attempts to pick an appropriate compression
algorithm for the dataset.
reader_func: Optional. A function to control how to read data from snapshot
shards.
shard_func: Optional. A function to control how to shard data when writing a
snapshot.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
"""Actual dataset transformation."""
return dataset.snapshot(
path=path,
compression=compression,
reader_func=reader_func,
shard_func=shard_func)
return _apply_fn
tensorflow/python/data/experimental/ops/take_while_ops.py
+38
View File
@@ -0,0 +1,38 @@
# Copyright 2019 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.
# ==============================================================================
"""take-while dataset transformation."""
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@deprecation.deprecated(None, "Use `tf.data.Dataset.take_while(...)")
@tf_export("data.experimental.take_while")
def take_while(predicate):
"""A transformation that stops dataset iteration based on a `predicate`.
Args:
predicate: A function that maps a nested structure of tensors (having shapes
and types defined by `self.output_shapes` and `self.output_types`) to a
scalar `tf.bool` tensor.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.take_while(predicate=predicate)
return _apply_fn
tensorflow/python/data/experimental/ops/testing.py
+198
View File
@@ -0,0 +1,198 @@
# 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 testing of tf.data."""
from google.protobuf import text_format
from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_experimental_dataset_ops
def assert_next(transformations):
"""A transformation that asserts which transformations happen next.
Transformations should be referred to by their base name, not including
version suffix. For example, use "Batch" instead of "BatchV2". "Batch" will
match any of "Batch", "BatchV1", "BatchV2", etc.
Args:
transformations: A `tf.string` vector `tf.Tensor` identifying the
transformations that are expected to happen next.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
return _AssertNextDataset(dataset, transformations)
return _apply_fn
def assert_prev(transformations):
r"""Asserts which transformations, with which attributes, happened previously.
Each transformation is represented as a tuple in the input.
The first element is the base op name of the transformation, not including
version suffix. For example, use "BatchDataset" instead of
"BatchDatasetV2". "BatchDataset" will match any of "BatchDataset",
"BatchDatasetV1", "BatchDatasetV2", etc.
The second element is a dict of attribute name-value pairs. Attributes
values must be of type bool, int, or string.
Example usage:
>>> dataset_ops.Dataset.from_tensors(0) \
... .map(lambda x: x) \
... .batch(1, deterministic=True, num_parallel_calls=8) \
... .assert_prev([("ParallelBatchDataset", {"deterministic": True}), \
... ("MapDataset", {})])
Args:
transformations: A list of tuples identifying the (required) transformation
name, with (optional) attribute name-value pairs, that are expected to
have happened previously.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
return _AssertPrevDataset(dataset, transformations)
return _apply_fn
def non_serializable():
"""A non-serializable identity transformation.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
return _NonSerializableDataset(dataset)
return _apply_fn
def sleep(sleep_microseconds):
"""Sleeps for `sleep_microseconds` before producing each input element.
Args:
sleep_microseconds: The number of microseconds to sleep before producing an
input element.
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return _SleepDataset(dataset, sleep_microseconds)
return _apply_fn
class _AssertNextDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that asserts which transformations happen next."""
def __init__(self, input_dataset, transformations):
"""See `assert_next()` for details."""
self._input_dataset = input_dataset
if transformations is None:
raise ValueError(
"Invalid `transformations`. `transformations` should not be empty.")
self._transformations = ops.convert_to_tensor(
transformations, dtype=dtypes.string, name="transformations")
variant_tensor = (
gen_experimental_dataset_ops.experimental_assert_next_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._transformations,
**self._flat_structure))
super(_AssertNextDataset, self).__init__(input_dataset, variant_tensor)
class _AssertPrevDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that asserts which transformations happened previously."""
def __init__(self, input_dataset, transformations):
"""See `assert_prev()` for details."""
self._input_dataset = input_dataset
if transformations is None:
raise ValueError("`transformations` cannot be empty")
def serialize_transformation(op_name, attributes):
proto = attr_value_pb2.NameAttrList(name=op_name)
if attributes is None or isinstance(attributes, set):
attributes = dict()
for (name, value) in attributes.items():
if isinstance(value, bool):
proto.attr[name].b = value
elif isinstance(value, int):
proto.attr[name].i = value
elif isinstance(value, str):
proto.attr[name].s = value.encode()
else:
raise ValueError(
f"attribute value type ({type(value)}) must be bool, int, or str")
return text_format.MessageToString(proto)
self._transformations = ops.convert_to_tensor(
[serialize_transformation(*x) for x in transformations],
dtype=dtypes.string,
name="transformations")
variant_tensor = (
gen_experimental_dataset_ops.assert_prev_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._transformations,
**self._flat_structure))
super(_AssertPrevDataset, self).__init__(input_dataset, variant_tensor)
class _NonSerializableDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that performs non-serializable identity transformation."""
def __init__(self, input_dataset):
"""See `non_serializable()` for details."""
self._input_dataset = input_dataset
variant_tensor = (
gen_experimental_dataset_ops.experimental_non_serializable_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
**self._flat_structure))
super(_NonSerializableDataset, self).__init__(input_dataset, variant_tensor)
class _SleepDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A `Dataset` that sleeps before producing each upstream element."""
def __init__(self, input_dataset, sleep_microseconds):
self._input_dataset = input_dataset
self._sleep_microseconds = sleep_microseconds
variant_tensor = gen_experimental_dataset_ops.sleep_dataset(
self._input_dataset._variant_tensor, # pylint: disable=protected-access
self._sleep_microseconds,
**self._flat_structure)
super(_SleepDataset, self).__init__(input_dataset, variant_tensor)
tensorflow/python/data/experimental/ops/unique.py
+43
View File
@@ -0,0 +1,43 @@
# Copyright 2017 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.
# ==============================================================================
"""Unique element dataset transformations."""
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@deprecation.deprecated(None, "Use `tf.data.Dataset.unique(...)")
@tf_export("data.experimental.unique")
def unique():
"""Creates a `Dataset` from another `Dataset`, discarding duplicates.
Use this transformation to produce a dataset that contains one instance of
each unique element in the input. For example:
```python
dataset = tf.data.Dataset.from_tensor_slices([1, 37, 2, 37, 2, 1])
# Using `unique()` will drop the duplicate elements.
dataset = dataset.apply(tf.data.experimental.unique()) # ==> { 1, 37, 2 }
```
Returns:
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
return dataset.unique()
return _apply_fn
tensorflow/python/data/experimental/ops/weighted_flat_map_op.py
+162
View File
@@ -0,0 +1,162 @@
# Copyright 2024 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.
# ==============================================================================
"""The implementation of `weighted_flat_map`."""
from typing import Optional, Sequence, Union
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest as tf_nest
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
def _weighted_flat_map( # pylint: disable=unused-private-name
input_datasets: Sequence[dataset_ops.DatasetV2],
weights: Optional[Sequence[Union[float, tensor.Tensor]]] = None,
name: Optional[str] = None) -> dataset_ops.DatasetV2:
"""A `Dataset` that fetches elements from `input_datasets` and flattens them.
This operation combines elements from multiple datasets into a flattened
dataset. Elements are read in proportion to the `weights` assigned to each
input dataset. All requested elements from a dataset are read before reading
the elements from the next dataset.
For example, suppose we have 2 datasets:
# TODO(wilsin): Make the following code testable after the API is released.
dataset1 = tf.data.Dataset.range(0, 10)
dataset2 = tf.data.Dataset.range(10, 20),
Suppose that we call `weighted_flat_map` from these 2 datasets with the
following weights:
dataset = tf.data.Dataset.weighted_flat_map([dataset1, dataset2], [0.5, 1.0])
Then, the outcome of the elements is:
# [0, 1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
Args:
input_datasets: A non-empty list of `tf.data.Dataset` objects with
compatible structure.
weights: (Optional.) A list or Tensor of `len(datasets)` non-zero
floating-point values where `weights[i]` represents the probability to
sample from `datasets[i]`, or a `tf.data.Dataset` object where each
element is such a list. Defaults to a uniform distribution across
`datasets`.
name: (Optional.) A name for the tf.data operation.
Returns:
A dataset that reads elements from all its inputs, reading the requested
elements from an input according to the weight before proceeding to the next
input. The number of elements read from an input is in proportion to its
weight given in `weights`.
Raises:
TypeError: if the `datasets` or `weights` arguments have the wrong type.
ValueError:
- if `input_datasets` has less than 2 datasets.
- if `weights` is specified and does not match the length of
`input_datasets`.
InvalidArgumentError:
- if any of the `input_datasets` has an unknown or infinite cardinality.
- if any of the `weights` has a value that is less than or equal to 0.0
"""
return _WeightedFlatMap(input_datasets, weights, name=name)
class _WeightedFlatMap(dataset_ops.DatasetV2):
"""A `Dataset` that maps a function over its input and flattens the result."""
def __init__(
self,
input_datasets: Sequence[dataset_ops.DatasetV2],
weights: Optional[Sequence[Union[float, tensor.Tensor]]] = None,
name: Optional[str] = None):
if not input_datasets:
raise ValueError("Invalid `datasets`. `datasets` should not be empty.")
if len(input_datasets) < 2:
raise ValueError(
"Invalid `datasets`. `datasets` should have at least two datasets.")
self._input_datasets = input_datasets
self._name = name
def common_supertype(a, b):
result = a.most_specific_common_supertype([b])
if result is None:
raise TypeError(f"No common supertype of {a} and {b}.")
return result
self._structure = input_datasets[0].element_spec
for dataset in input_datasets[1:]:
try:
self._structure = tf_nest.map_structure(
common_supertype, self._structure, dataset.element_spec
)
except (TypeError, ValueError) as e:
raise TypeError(
"Incompatible dataset elements:\n"
f" {input_datasets[0].element_spec} vs. "
f" {dataset.element_spec}"
) from e
if weights is None:
weights = [1.0] * len(input_datasets)
else:
if isinstance(weights, tensor.Tensor):
if not weights.shape.is_compatible_with([len(input_datasets)]):
raise ValueError(
"Invalid `weights`. The shape of `weights` "
"should be compatible with `[len(datasets)]` "
f"but is {weights.shape}."
)
else:
if len(input_datasets) != len(weights):
raise ValueError(
"Invalid `weights`. `weights` should have the "
"same length as `datasets` but got "
f"`len(weights)={len(weights)}` vs. "
f"`len(datasets)={len(input_datasets)}`."
)
weights = [
ops.convert_to_tensor(w, preferred_dtype=dtypes.float64)
for w in weights
]
for weight in weights:
if weight.dtype not in (dtypes.float32, dtypes.float64):
raise TypeError(
"Invalid `weights`. `weights` type must be either "
"`tf.float32` or `tf.float64` but is "
f"{weight.dtype}."
)
# pylint: disable=protected-access
variant_tensor = ged_ops.weighted_flat_map_dataset(
[dataset._variant_tensor for dataset in self._input_datasets],
weights,
**self._common_args,
)
super().__init__(variant_tensor)
def _inputs(self):
return self._input_datasets
@property
def element_spec(self):
return self._structure
def _transformation_name(self):
return "Dataset.weighted_flat_map()"
tensorflow/python/data/experimental/ops/writers.py
+126
View File
@@ -0,0 +1,126 @@
# 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.
# ==============================================================================
"""Python wrappers for tf.data writers."""
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import convert
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.types import data as data_types
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@tf_export("data.experimental.TFRecordWriter")
@deprecation.deprecated(
None, "To write TFRecords to disk, use `tf.io.TFRecordWriter`. To save "
"and load the contents of a dataset, use `tf.data.experimental.save` "
"and `tf.data.experimental.load`")
class TFRecordWriter:
"""Writes a dataset to a TFRecord file.
The elements of the dataset must be scalar strings. To serialize dataset
elements as strings, you can use the `tf.io.serialize_tensor` function.
```python
dataset = tf.data.Dataset.range(3)
dataset = dataset.map(tf.io.serialize_tensor)
writer = tf.data.experimental.TFRecordWriter("/path/to/file.tfrecord")
writer.write(dataset)
```
To read back the elements, use `TFRecordDataset`.
```python
dataset = tf.data.TFRecordDataset("/path/to/file.tfrecord")
dataset = dataset.map(lambda x: tf.io.parse_tensor(x, tf.int64))
```
To shard a `dataset` across multiple TFRecord files:
```python
dataset = ... # dataset to be written
def reduce_func(key, dataset):
filename = tf.strings.join([PATH_PREFIX, tf.strings.as_string(key)])
writer = tf.data.experimental.TFRecordWriter(filename)
writer.write(dataset.map(lambda _, x: x))
return tf.data.Dataset.from_tensors(filename)
dataset = dataset.enumerate()
dataset = dataset.apply(tf.data.experimental.group_by_window(
lambda i, _: i % NUM_SHARDS, reduce_func, tf.int64.max
))
# Iterate through the dataset to trigger data writing.
for _ in dataset:
pass
```
"""
def __init__(self, filename, compression_type=None):
"""Initializes a `TFRecordWriter`.
Args:
filename: a string path indicating where to write the TFRecord data.
compression_type: (Optional.) a string indicating what type of compression
to use when writing the file. See `tf.io.TFRecordCompressionType` for
what types of compression are available. Defaults to `None`.
"""
self._filename = ops.convert_to_tensor(
filename, dtypes.string, name="filename")
self._compression_type = convert.optional_param_to_tensor(
"compression_type",
compression_type,
argument_default="",
argument_dtype=dtypes.string)
def write(self, dataset):
"""Writes a dataset to a TFRecord file.
An operation that writes the content of the specified dataset to the file
specified in the constructor.
If the file exists, it will be overwritten.
Args:
dataset: a `tf.data.Dataset` whose elements are to be written to a file
Returns:
In graph mode, this returns an operation which when executed performs the
write. In eager mode, the write is performed by the method itself and
there is no return value.
Raises
TypeError: if `dataset` is not a `tf.data.Dataset`.
TypeError: if the elements produced by the dataset are not scalar strings.
"""
if not isinstance(dataset, data_types.DatasetV2):
raise TypeError(
f"Invalid `dataset.` Expected a `tf.data.Dataset` object but got "
f"{type(dataset)}."
)
if not dataset_ops.get_structure(dataset).is_compatible_with(
tensor_spec.TensorSpec([], dtypes.string)):
raise TypeError(
f"Invalid `dataset`. Expected a`dataset` that produces scalar "
f"`tf.string` elements, but got a dataset which produces elements "
f"with shapes {dataset_ops.get_legacy_output_shapes(dataset)} and "
f"types {dataset_ops.get_legacy_output_types(dataset)}.")
# pylint: disable=protected-access
dataset = dataset._apply_debug_options()
return gen_experimental_dataset_ops.dataset_to_tf_record(
dataset._variant_tensor, self._filename, self._compression_type)