# 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 Iterators.""" import abc import threading import warnings import six from tensorflow.python.data.ops import optional_ops from tensorflow.python.data.ops import options as options_lib from tensorflow.python.data.util import nest from tensorflow.python.data.util import structure from tensorflow.python.eager import context from tensorflow.python.framework import composite_tensor from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import tensor_spec from tensorflow.python.framework import type_spec from tensorflow.python.ops import gen_dataset_ops from tensorflow.python.trackable import base as trackable from tensorflow.python.training.saver import BaseSaverBuilder from tensorflow.python.util import _pywrap_utils from tensorflow.python.util import deprecation from tensorflow.python.util import lazy_loader from tensorflow.python.util.compat import collections_abc from tensorflow.python.util.tf_export import tf_export # NOTE(mrry): It is legitimate to call `Iterator.get_next()` multiple # times, e.g. when you are distributing different elements to multiple # devices in a single step. However, a common pitfall arises when # users call `Iterator.get_next()` in each iteration of their training # loop. `Iterator.get_next()` adds ops to the graph, and executing # each op allocates resources (including threads); as a consequence, # invoking it in every iteration of a training loop causes slowdown # and eventual resource exhaustion. To guard against this outcome, we # log a warning when the number of uses crosses a threshold of suspicion. GET_NEXT_CALL_WARNING_THRESHOLD = 32 GET_NEXT_CALL_WARNING_MESSAGE = ( "An unusually high number of `Iterator.get_next()` calls was detected. " "This often indicates that `Iterator.get_next()` is being called inside " "a training loop, which will cause gradual slowdown and eventual resource " "exhaustion. If this is the case, restructure your code to call " "`next_element = iterator.get_next()` once outside the loop, and use " "`next_element` as the input to some computation that is invoked inside " "the loop.") # NOTE(jsimsa): Threshold used as a heuristic to check for infinite loop during # tf.function tracing. GET_NEXT_CALL_ERROR_THRESHOLD = 32 GET_NEXT_CALL_ERROR_MESSAGE = ( "An unusually high number of `tf.data.Iterator.get_next()` calls was " "detected. This suggests that the `for elem in dataset: ...` idiom is used " "within tf.function with AutoGraph disabled. This idiom is only supported " "when AutoGraph is enabled.") # Collection of all IteratorResources in the `Graph`. GLOBAL_ITERATORS = "iterators" autograph_ctx = lazy_loader.LazyLoader( "autograph_ctx", globals(), "tensorflow.python.autograph.core.ag_ctx") def _device_stack_is_empty(): if context.executing_eagerly(): return context.context().device_name is None # pylint: disable=protected-access device_stack = ops.get_default_graph()._device_functions_outer_to_inner # pylint: enable=protected-access return not bool(device_stack) @tf_export(v1=["data.Iterator"]) class Iterator(trackable.Trackable): """Represents the state of iterating through a `Dataset`.""" def __init__(self, iterator_resource, initializer, output_types, output_shapes, output_classes): """Creates a new iterator from the given iterator resource. Note: Most users will not call this initializer directly, and will instead use `Dataset.make_initializable_iterator()` or `Dataset.make_one_shot_iterator()`. Args: iterator_resource: A `tf.resource` scalar `tf.Tensor` representing the iterator. initializer: A `tf.Operation` that should be run to initialize this iterator. output_types: A (nested) structure of `tf.DType` objects corresponding to each component of an element of this iterator. output_shapes: A (nested) structure of `tf.TensorShape` objects corresponding to each component of an element of this iterator. output_classes: A (nested) structure of Python `type` objects corresponding to each component of an element of this iterator. Raises: TypeError: If `output_types`, `output_shapes`, or `output_classes` is not specified. """ self._iterator_resource = iterator_resource self._initializer = initializer if (output_types is None or output_shapes is None or output_classes is None): raise ValueError( "All of `output_types`, `output_shapes`, and `output_classes` " "must be specified to create an iterator. Got " f"`output_types` = {output_types!r}, " f"`output_shapes` = {output_shapes!r}, " f"`output_classes` = {output_classes!r}.") self._element_spec = structure.convert_legacy_structure( output_types, output_shapes, output_classes) self._flat_tensor_shapes = structure.get_flat_tensor_shapes( self._element_spec) self._flat_tensor_types = structure.get_flat_tensor_types( self._element_spec) self._string_handle = gen_dataset_ops.iterator_to_string_handle( self._iterator_resource) self._get_next_call_count = 0 ops.add_to_collection(GLOBAL_ITERATORS, self._iterator_resource) @staticmethod def from_structure(output_types, output_shapes=None, shared_name=None, output_classes=None): """Creates a new, uninitialized `Iterator` with the given structure. This iterator-constructing method can be used to create an iterator that is reusable with many different datasets. The returned iterator is not bound to a particular dataset, and it has no `initializer`. To initialize the iterator, run the operation returned by `Iterator.make_initializer(dataset)`. The following is an example ```python iterator = Iterator.from_structure(tf.int64, tf.TensorShape([])) dataset_range = Dataset.range(10) range_initializer = iterator.make_initializer(dataset_range) dataset_evens = dataset_range.filter(lambda x: x % 2 == 0) evens_initializer = iterator.make_initializer(dataset_evens) # Define a model based on the iterator; in this example, the model_fn # is expected to take scalar tf.int64 Tensors as input (see # the definition of 'iterator' above). prediction, loss = model_fn(iterator.get_next()) # Train for `num_epochs`, where for each epoch, we first iterate over # dataset_range, and then iterate over dataset_evens. for _ in range(num_epochs): # Initialize the iterator to `dataset_range` sess.run(range_initializer) while True: try: pred, loss_val = sess.run([prediction, loss]) except tf.errors.OutOfRangeError: break # Initialize the iterator to `dataset_evens` sess.run(evens_initializer) while True: try: pred, loss_val = sess.run([prediction, loss]) except tf.errors.OutOfRangeError: break ``` Args: output_types: A (nested) structure of `tf.DType` objects corresponding to each component of an element of this dataset. output_shapes: (Optional.) A (nested) structure of `tf.TensorShape` objects corresponding to each component of an element of this dataset. If omitted, each component will have an unconstrainted shape. shared_name: (Optional.) If non-empty, this iterator will be shared under the given name across multiple sessions that share the same devices (e.g. when using a remote server). output_classes: (Optional.) A (nested) structure of Python `type` objects corresponding to each component of an element of this iterator. If omitted, each component is assumed to be of type `tf.Tensor`. Returns: An `Iterator`. Raises: TypeError: If the structures of `output_shapes` and `output_types` are not the same. """ output_types = nest.map_structure(dtypes.as_dtype, output_types) if output_shapes is None: output_shapes = nest.map_structure( lambda _: tensor_shape.TensorShape(None), output_types) else: output_shapes = nest.map_structure_up_to(output_types, tensor_shape.as_shape, output_shapes) if output_classes is None: output_classes = nest.map_structure(lambda _: ops.Tensor, output_types) nest.assert_same_structure(output_types, output_shapes) output_structure = structure.convert_legacy_structure( output_types, output_shapes, output_classes) if shared_name is None: shared_name = "" iterator_resource = gen_dataset_ops.iterator_v2( container="", shared_name=shared_name, output_types=structure.get_flat_tensor_types(output_structure), output_shapes=structure.get_flat_tensor_shapes( output_structure)) return Iterator(iterator_resource, None, output_types, output_shapes, output_classes) @staticmethod def from_string_handle(string_handle, output_types, output_shapes=None, output_classes=None): """Creates a new, uninitialized `Iterator` based on the given handle. This method allows you to define a "feedable" iterator where you can choose between concrete iterators by feeding a value in a `tf.Session.run` call. In that case, `string_handle` would be a `tf.compat.v1.placeholder`, and you would feed it with the value of `tf.data.Iterator.string_handle` in each step. For example, if you had two iterators that marked the current position in a training dataset and a test dataset, you could choose which to use in each step as follows: ```python train_iterator = tf.data.Dataset(...).make_one_shot_iterator() train_iterator_handle = sess.run(train_iterator.string_handle()) test_iterator = tf.data.Dataset(...).make_one_shot_iterator() test_iterator_handle = sess.run(test_iterator.string_handle()) handle = tf.compat.v1.placeholder(tf.string, shape=[]) iterator = tf.data.Iterator.from_string_handle( handle, train_iterator.output_types) next_element = iterator.get_next() loss = f(next_element) train_loss = sess.run(loss, feed_dict={handle: train_iterator_handle}) test_loss = sess.run(loss, feed_dict={handle: test_iterator_handle}) ``` Args: string_handle: A scalar `tf.Tensor` of type `tf.string` that evaluates to a handle produced by the `Iterator.string_handle()` method. output_types: A (nested) structure of `tf.DType` objects corresponding to each component of an element of this dataset. output_shapes: (Optional.) A (nested) structure of `tf.TensorShape` objects corresponding to each component of an element of this dataset. If omitted, each component will have an unconstrainted shape. output_classes: (Optional.) A (nested) structure of Python `type` objects corresponding to each component of an element of this iterator. If omitted, each component is assumed to be of type `tf.Tensor`. Returns: An `Iterator`. """ output_types = nest.map_structure(dtypes.as_dtype, output_types) if output_shapes is None: output_shapes = nest.map_structure( lambda _: tensor_shape.TensorShape(None), output_types) else: output_shapes = nest.map_structure_up_to(output_types, tensor_shape.as_shape, output_shapes) if output_classes is None: output_classes = nest.map_structure(lambda _: ops.Tensor, output_types) nest.assert_same_structure(output_types, output_shapes) output_structure = structure.convert_legacy_structure( output_types, output_shapes, output_classes) string_handle = ops.convert_to_tensor(string_handle, dtype=dtypes.string) iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2( string_handle, output_types=structure.get_flat_tensor_types(output_structure), output_shapes=structure.get_flat_tensor_shapes(output_structure)) return Iterator(iterator_resource, None, output_types, output_shapes, output_classes) @property def initializer(self): """A `tf.Operation` that should be run to initialize this iterator. Returns: A `tf.Operation` that should be run to initialize this iterator Raises: ValueError: If this iterator initializes itself automatically. """ if self._initializer is not None: return self._initializer else: # TODO(mrry): Consider whether one-shot iterators should have # initializers that simply reset their state to the beginning. raise ValueError( "The iterator does not have an initializer. This means it was likely " "created using `tf.data.Dataset.make_one_shot_iterator()`. For an " "initializable iterator, use " "`tf.data.Dataset.make_initializable_iterator()` instead.") def make_initializer(self, dataset, name=None): """Returns a `tf.Operation` that initializes this iterator on `dataset`. Args: dataset: A `Dataset` whose `element_spec` if compatible with this iterator. name: (Optional.) A name for the created operation. Returns: A `tf.Operation` that can be run to initialize this iterator on the given `dataset`. Raises: TypeError: If `dataset` and this iterator do not have a compatible `element_spec`. """ with ops.name_scope(name, "make_initializer") as name: # NOTE(mrry): Cannot depend on `dataset_ops.get_legacy_output*()` due # to that creating a circular dependency. # pylint: disable=protected-access dataset_output_types = nest.map_structure( lambda component_spec: component_spec._to_legacy_output_types(), dataset.element_spec) dataset_output_shapes = nest.map_structure( lambda component_spec: component_spec._to_legacy_output_shapes(), dataset.element_spec) dataset_output_classes = nest.map_structure( lambda component_spec: component_spec._to_legacy_output_classes(), dataset.element_spec) # pylint: enable=protected-access nest.assert_same_structure(self.output_types, dataset_output_types) nest.assert_same_structure(self.output_shapes, dataset_output_shapes) for iterator_class, dataset_class in zip( nest.flatten(self.output_classes), nest.flatten(dataset_output_classes)): if iterator_class is not dataset_class: raise TypeError( f"Expected output classes {self.output_classes!r} but got " f"dataset with output classes {dataset_output_classes!r}.") for iterator_dtype, dataset_dtype in zip( nest.flatten(self.output_types), nest.flatten(dataset_output_types)): if iterator_dtype != dataset_dtype: raise TypeError( f"Expected output types {self.output_types!r} but got dataset " f"with output types {dataset_output_types!r}.") for iterator_shape, dataset_shape in zip( nest.flatten(self.output_shapes), nest.flatten( dataset_output_shapes)): if not iterator_shape.is_compatible_with(dataset_shape): raise TypeError( f"Expected output shapes compatible with {self.output_shapes!r} " f"but got dataset with output shapes {dataset_output_shapes!r}.") # TODO(b/169442955): Investigate the need for this colocation constraint. with ops.colocate_with(self._iterator_resource): # pylint: disable=protected-access return gen_dataset_ops.make_iterator( dataset._variant_tensor, self._iterator_resource, name=name) def get_next(self, name=None): """Returns the next element. In graph mode, you should typically call this method *once* and use its result as the input to another computation. A typical loop will then call `tf.Session.run` on the result of that computation. The loop will terminate when the `Iterator.get_next()` operation raises `tf.errors.OutOfRangeError`. The following skeleton shows how to use this method when building a training loop: ```python dataset = ... # A `tf.data.Dataset` object. iterator = dataset.make_initializable_iterator() next_element = iterator.get_next() # Build a TensorFlow graph that does something with each element. loss = model_function(next_element) optimizer = ... # A `tf.compat.v1.train.Optimizer` object. train_op = optimizer.minimize(loss) with tf.compat.v1.Session() as sess: try: while True: sess.run(train_op) except tf.errors.OutOfRangeError: pass ``` NOTE: It is legitimate to call `Iterator.get_next()` multiple times, e.g. when you are distributing different elements to multiple devices in a single step. However, a common pitfall arises when users call `Iterator.get_next()` in each iteration of their training loop. `Iterator.get_next()` adds ops to the graph, and executing each op allocates resources (including threads); as a consequence, invoking it in every iteration of a training loop causes slowdown and eventual resource exhaustion. To guard against this outcome, we log a warning when the number of uses crosses a fixed threshold of suspiciousness. Args: name: (Optional.) A name for the created operation. Returns: A (nested) structure of values matching `tf.data.Iterator.element_spec`. """ self._get_next_call_count += 1 if self._get_next_call_count > GET_NEXT_CALL_WARNING_THRESHOLD: warnings.warn(GET_NEXT_CALL_WARNING_MESSAGE) # TODO(b/169442955): Investigate the need for this colocation constraint. with ops.colocate_with(self._iterator_resource): # pylint: disable=protected-access flat_ret = gen_dataset_ops.iterator_get_next( self._iterator_resource, output_types=self._flat_tensor_types, output_shapes=self._flat_tensor_shapes, name=name) return structure.from_tensor_list(self._element_spec, flat_ret) def get_next_as_optional(self): # TODO(b/169442955): Investigate the need for this colocation constraint. with ops.colocate_with(self._iterator_resource): # pylint: disable=protected-access return optional_ops._OptionalImpl( gen_dataset_ops.iterator_get_next_as_optional( self._iterator_resource, output_types=structure.get_flat_tensor_types(self.element_spec), output_shapes=structure.get_flat_tensor_shapes( self.element_spec)), self.element_spec) def string_handle(self, name=None): """Returns a string-valued `tf.Tensor` that represents this iterator. Args: name: (Optional.) A name for the created operation. Returns: A scalar `tf.Tensor` of type `tf.string`. """ if name is None: return self._string_handle else: return gen_dataset_ops.iterator_to_string_handle( self._iterator_resource, name=name) @property @deprecation.deprecated( None, "Use `tf.compat.v1.data.get_output_classes(iterator)`.") def output_classes(self): """Returns the class of each component of an element of this iterator. The expected values are `tf.Tensor` and `tf.sparse.SparseTensor`. Returns: A (nested) structure of Python `type` objects corresponding to each component of an element of this dataset. """ return nest.map_structure( lambda component_spec: component_spec._to_legacy_output_classes(), # pylint: disable=protected-access self._element_spec) @property @deprecation.deprecated( None, "Use `tf.compat.v1.data.get_output_shapes(iterator)`.") def output_shapes(self): """Returns the shape of each component of an element of this iterator. Returns: A (nested) structure of `tf.TensorShape` objects corresponding to each component of an element of this dataset. """ return nest.map_structure( lambda component_spec: component_spec._to_legacy_output_shapes(), # pylint: disable=protected-access self._element_spec) @property @deprecation.deprecated( None, "Use `tf.compat.v1.data.get_output_types(iterator)`.") def output_types(self): """Returns the type of each component of an element of this iterator. Returns: A (nested) structure of `tf.DType` objects corresponding to each component of an element of this dataset. """ return nest.map_structure( lambda component_spec: component_spec._to_legacy_output_types(), # pylint: disable=protected-access self._element_spec) @property def element_spec(self): """The type specification of an element of this iterator. For more information, read [this guide](https://www.tensorflow.org/guide/data#dataset_structure). Returns: A (nested) structure of `tf.TypeSpec` objects matching the structure of an element of this iterator and specifying the type of individual components. """ return self._element_spec def _gather_saveables_for_checkpoint(self): def _saveable_factory(name): return _IteratorSaveable(self._iterator_resource, name) return {"ITERATOR": _saveable_factory} _uid_counter = 0 _uid_lock = threading.Lock() def _generate_shared_name(prefix): with _uid_lock: global _uid_counter uid = _uid_counter _uid_counter += 1 return "{}{}".format(prefix, uid) @tf_export("data.Iterator", v1=[]) @six.add_metaclass(abc.ABCMeta) class IteratorBase(collections_abc.Iterator, trackable.Trackable, composite_tensor.CompositeTensor): """Represents an iterator of a `tf.data.Dataset`. `tf.data.Iterator` is the primary mechanism for enumerating elements of a `tf.data.Dataset`. It supports the Python Iterator protocol, which means it can be iterated over using a for-loop: >>> dataset = tf.data.Dataset.range(2) >>> for element in dataset: ... print(element) tf.Tensor(0, shape=(), dtype=int64) tf.Tensor(1, shape=(), dtype=int64) or by fetching individual elements explicitly via `get_next()`: >>> dataset = tf.data.Dataset.range(2) >>> iterator = iter(dataset) >>> print(iterator.get_next()) tf.Tensor(0, shape=(), dtype=int64) >>> print(iterator.get_next()) tf.Tensor(1, shape=(), dtype=int64) In addition, non-raising iteration is supported via `get_next_as_optional()`, which returns the next element (if available) wrapped in a `tf.experimental.Optional`. >>> dataset = tf.data.Dataset.from_tensors(42) >>> iterator = iter(dataset) >>> optional = iterator.get_next_as_optional() >>> print(optional.has_value()) tf.Tensor(True, shape=(), dtype=bool) >>> optional = iterator.get_next_as_optional() >>> print(optional.has_value()) tf.Tensor(False, shape=(), dtype=bool) """ @abc.abstractproperty def element_spec(self): """The type specification of an element of this iterator. >>> dataset = tf.data.Dataset.from_tensors(42) >>> iterator = iter(dataset) >>> iterator.element_spec tf.TensorSpec(shape=(), dtype=tf.int32, name=None) For more information, read [this guide](https://www.tensorflow.org/guide/data#dataset_structure). Returns: A (nested) structure of `tf.TypeSpec` objects matching the structure of an element of this iterator, specifying the type of individual components. """ raise NotImplementedError("Iterator.element_spec") @abc.abstractmethod def get_next(self): """Returns the next element. >>> dataset = tf.data.Dataset.from_tensors(42) >>> iterator = iter(dataset) >>> print(iterator.get_next()) tf.Tensor(42, shape=(), dtype=int32) Returns: A (nested) structure of values matching `tf.data.Iterator.element_spec`. Raises: `tf.errors.OutOfRangeError`: If the end of the iterator has been reached. """ raise NotImplementedError("Iterator.get_next()") @abc.abstractmethod def get_next_as_optional(self): """Returns the next element wrapped in `tf.experimental.Optional`. If the iterator has reached the end of the sequence, the returned `tf.experimental.Optional` will have no value. >>> dataset = tf.data.Dataset.from_tensors(42) >>> iterator = iter(dataset) >>> optional = iterator.get_next_as_optional() >>> print(optional.has_value()) tf.Tensor(True, shape=(), dtype=bool) >>> print(optional.get_value()) tf.Tensor(42, shape=(), dtype=int32) >>> optional = iterator.get_next_as_optional() >>> print(optional.has_value()) tf.Tensor(False, shape=(), dtype=bool) Returns: A `tf.experimental.Optional` object representing the next element. """ raise NotImplementedError("Iterator.get_next_as_optional()") class OwnedIterator(IteratorBase): """An iterator producing tf.Tensor objects from a tf.data.Dataset. The iterator resource created through `OwnedIterator` is owned by the Python object and the life time of the underlying resource is tied to the life time of the `OwnedIterator` object. This makes `OwnedIterator` appropriate for use in eager mode and inside of tf.functions. """ def __init__(self, dataset=None, components=None, element_spec=None): """Creates a new iterator from the given dataset. If `dataset` is not specified, the iterator will be created from the given tensor components and element structure. In particular, the alternative for constructing the iterator is used when the iterator is reconstructed from it `CompositeTensor` representation. Args: dataset: A `tf.data.Dataset` object. components: Tensor components to construct the iterator from. element_spec: A (nested) structure of `TypeSpec` objects that represents the type specification of elements of the iterator. Raises: ValueError: If `dataset` is not provided and either `components` or `element_spec` is not provided. Or `dataset` is provided and either `components` and `element_spec` is provided. """ super(OwnedIterator, self).__init__() if dataset is None: if (components is None or element_spec is None): raise ValueError( "When `dataset` is not provided, both `components` and " "`element_spec` must be specified.") # pylint: disable=protected-access self._element_spec = element_spec self._flat_output_types = structure.get_flat_tensor_types( self._element_spec) self._flat_output_shapes = structure.get_flat_tensor_shapes( self._element_spec) self._iterator_resource, = components else: if (components is not None or element_spec is not None): raise ValueError( "When `dataset` is provided, `element_spec` and `components` must " "not be specified.") self._create_iterator(dataset) self._get_next_call_count = 0 def _create_iterator(self, dataset): # pylint: disable=protected-access dataset = dataset._apply_debug_options() # Store dataset reference to ensure that dataset is alive when this iterator # is being used. For example, `tf.data.Dataset.from_generator` registers # a few py_funcs that are needed in `self._next_internal`. If the dataset # is deleted, this iterator crashes on `self.__next__(...)` call. self._dataset = dataset ds_variant = dataset._variant_tensor self._element_spec = dataset.element_spec self._flat_output_types = structure.get_flat_tensor_types( self._element_spec) self._flat_output_shapes = structure.get_flat_tensor_shapes( self._element_spec) with ops.colocate_with(ds_variant): self._iterator_resource = ( gen_dataset_ops.anonymous_iterator_v3( output_types=self._flat_output_types, output_shapes=self._flat_output_shapes)) gen_dataset_ops.make_iterator(ds_variant, self._iterator_resource) def __iter__(self): return self def next(self): # For Python 2 compatibility return self.__next__() def _next_internal(self): autograph_status = autograph_ctx.control_status_ctx().status autograph_disabled = autograph_status == autograph_ctx.Status.DISABLED if not context.executing_eagerly() and autograph_disabled: self._get_next_call_count += 1 if self._get_next_call_count > GET_NEXT_CALL_ERROR_THRESHOLD: raise ValueError(GET_NEXT_CALL_ERROR_MESSAGE) if not context.executing_eagerly(): # TODO(b/169442955): Investigate the need for this colocation constraint. with ops.colocate_with(self._iterator_resource): ret = gen_dataset_ops.iterator_get_next( self._iterator_resource, output_types=self._flat_output_types, output_shapes=self._flat_output_shapes) return structure.from_compatible_tensor_list(self._element_spec, ret) # TODO(b/77291417): This runs in sync mode as iterators use an error status # to communicate that there is no more data to iterate over. with context.execution_mode(context.SYNC): ret = gen_dataset_ops.iterator_get_next( self._iterator_resource, output_types=self._flat_output_types, output_shapes=self._flat_output_shapes) try: # Fast path for the case `self._structure` is not a nested structure. return self._element_spec._from_compatible_tensor_list(ret) # pylint: disable=protected-access except AttributeError: return structure.from_compatible_tensor_list(self._element_spec, ret) @property def _type_spec(self): return IteratorSpec(self.element_spec) def __next__(self): try: return self._next_internal() except errors.OutOfRangeError: raise StopIteration @property @deprecation.deprecated( None, "Use `tf.compat.v1.data.get_output_classes(iterator)`.") def output_classes(self): """Returns the class of each component of an element of this iterator. The expected values are `tf.Tensor` and `tf.sparse.SparseTensor`. Returns: A (nested) structure of Python `type` objects corresponding to each component of an element of this dataset. """ return nest.map_structure( lambda component_spec: component_spec._to_legacy_output_classes(), # pylint: disable=protected-access self._element_spec) @property @deprecation.deprecated( None, "Use `tf.compat.v1.data.get_output_shapes(iterator)`.") def output_shapes(self): """Returns the shape of each component of an element of this iterator. Returns: A (nested) structure of `tf.TensorShape` objects corresponding to each component of an element of this dataset. """ return nest.map_structure( lambda component_spec: component_spec._to_legacy_output_shapes(), # pylint: disable=protected-access self._element_spec) @property @deprecation.deprecated( None, "Use `tf.compat.v1.data.get_output_types(iterator)`.") def output_types(self): """Returns the type of each component of an element of this iterator. Returns: A (nested) structure of `tf.DType` objects corresponding to each component of an element of this dataset. """ return nest.map_structure( lambda component_spec: component_spec._to_legacy_output_types(), # pylint: disable=protected-access self._element_spec) @property def element_spec(self): return self._element_spec def get_next(self): return self._next_internal() def get_next_as_optional(self): # TODO(b/169442955): Investigate the need for this colocation constraint. with ops.colocate_with(self._iterator_resource): # pylint: disable=protected-access return optional_ops._OptionalImpl( gen_dataset_ops.iterator_get_next_as_optional( self._iterator_resource, output_types=structure.get_flat_tensor_types(self.element_spec), output_shapes=structure.get_flat_tensor_shapes( self.element_spec)), self.element_spec) def _gather_saveables_for_checkpoint(self): def _saveable_factory(name): """Returns a SaveableObject for serialization/deserialization.""" policy = None if self._dataset: policy = self._dataset.options().experimental_external_state_policy if policy: return _IteratorSaveable( self._iterator_resource, name, external_state_policy=policy) else: return _IteratorSaveable(self._iterator_resource, name) return {"ITERATOR": _saveable_factory} def __tf_tracing_type__(self, signature_context): return signature_context.make_reference_type(self._type_spec, self._iterator_resource._id) # pylint:disable=protected-access @tf_export("data.IteratorSpec", v1=[]) class IteratorSpec(type_spec.TypeSpec): """Type specification for `tf.data.Iterator`. For instance, `tf.data.IteratorSpec` can be used to define a tf.function that takes `tf.data.Iterator` as an input argument: >>> @tf.function(input_signature=[tf.data.IteratorSpec( ... tf.TensorSpec(shape=(), dtype=tf.int32, name=None))]) ... def square(iterator): ... x = iterator.get_next() ... return x * x >>> dataset = tf.data.Dataset.from_tensors(5) >>> iterator = iter(dataset) >>> print(square(iterator)) tf.Tensor(25, shape=(), dtype=int32) Attributes: element_spec: A (nested) structure of `tf.TypeSpec` objects that represents the type specification of the iterator elements. """ __slots__ = ["_element_spec"] def __init__(self, element_spec): self._element_spec = element_spec @property def value_type(self): return OwnedIterator def _serialize(self): return (self._element_spec,) @property def _component_specs(self): return (tensor_spec.TensorSpec([], dtypes.resource),) def _to_components(self, value): return (value._iterator_resource,) # pylint: disable=protected-access def _from_components(self, components): return OwnedIterator( dataset=None, components=components, element_spec=self._element_spec) @staticmethod def from_value(value): return IteratorSpec(value.element_spec) # pylint: disable=protected-access def __tf_tracing_type__(self, signature_context): # TODO(b/202772221): Validate and enforce this assumption of uniqueness per # spec instance. return signature_context.make_reference_type(self, id(self)) # TODO(b/71645805): Expose trackable stateful objects from dataset. class _IteratorSaveable(BaseSaverBuilder.SaveableObject): """SaveableObject for saving/restoring iterator state.""" def __init__( self, iterator_resource, name, external_state_policy=options_lib.ExternalStatePolicy.FAIL): serialized_iterator = gen_dataset_ops.serialize_iterator( iterator_resource, external_state_policy=external_state_policy.value) specs = [ BaseSaverBuilder.SaveSpec( serialized_iterator, "", name + "_STATE", device=iterator_resource.device) ] super(_IteratorSaveable, self).__init__(iterator_resource, specs, name) def restore(self, restored_tensors, restored_shapes): with ops.colocate_with(self.op): return gen_dataset_ops.deserialize_iterator(self.op, restored_tensors[0]) @deprecation.deprecated( None, "Use `tf.data.Iterator.get_next_as_optional()` instead.") @tf_export("data.experimental.get_next_as_optional") def get_next_as_optional(iterator): """Returns a `tf.experimental.Optional` with the next element of the iterator. If the iterator has reached the end of the sequence, the returned `tf.experimental.Optional` will have no value. Args: iterator: A `tf.data.Iterator`. Returns: A `tf.experimental.Optional` object which either contains the next element of the iterator (if it exists) or no value. """ return iterator.get_next_as_optional() _pywrap_utils.RegisterType("OwnedIterator", OwnedIterator)