def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.contrib.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Args:
dataset: A `tf.contrib.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.in_eager_mode():
raise RuntimeError(
"{} objects only make sense when eager execution is enabled".
format(type(self)))
ds_variant = dataset.make_dataset_resource()
self._output_types = dataset.output_types
self._flat_output_types = nest.flatten(dataset.output_types)
self._flat_output_shapes = nest.flatten(dataset.output_shapes)
self._resource = gen_dataset_ops.iterator(
container="",
shared_name=_iterator_shared_name(),
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
def _create_iterator(self, dataset):
# pylint: disable=protected-access
dataset = dataset._apply_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, self._deleter = (
gen_dataset_ops.anonymous_iterator_v2(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes))
gen_dataset_ops.make_iterator(ds_variant, self._iterator_resource)
# Delete the resource when this object is deleted
self._resource_deleter = IteratorResourceDeleter(
handle=self._iterator_resource,
device=self._device,
deleter=self._deleter)
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.contrib.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Args:
dataset: A `tf.contrib.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.in_eager_mode():
raise RuntimeError(
"{} objects only make sense when eager execution is enabled".format(
type(self)))
ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access
self._output_types = dataset.output_types
self._flat_output_types = nest.flatten(dataset.output_types)
self._flat_output_shapes = nest.flatten(dataset.output_shapes)
self._resource = gen_dataset_ops.iterator(
container="",
shared_name=_iterator_shared_name(),
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
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))
if not context.executing_eagerly():
# Add full type information to the graph so host memory types inside
# variants stay on CPU, e.g, ragged string tensors.
# TODO(b/224776031) Remove this when AnonymousIterateV3 can use
# (reverse) type inference and all other ops that are needed to
# provide type information to the AnonymousIterateV3 also support
# type inference (esp. cross-function type inference) instead of
# setting the full type information manually.
fulltype = type_utils.iterator_full_type_from_spec(
self._element_spec)
# fulltype is PRODUCT[ITERATOR[PRODUCT[...]]]
assert len(fulltype.args[0].args[0].args) == len(
self._flat_output_types)
self._iterator_resource.op.experimental_set_type(fulltype)
gen_dataset_ops.make_iterator(ds_variant, self._iterator_resource)
def _eager_reset(self):
"""Resets the MultiDeviceIterator in eager mode."""
if not context.executing_eagerly():
raise ValueError("Eager reset is only supported in eager mode.")
# pylint: disable=protected-access
self._incarnation_id = gen_dataset_ops.multi_device_iterator_init(
self._dataset._variant_tensor,
self._multi_device_iterator_resource,
max_buffer_size=self._max_buffer_size)
for i, device in enumerate(self._devices):
with ops.device(device):
ds = self._create_device_dataset(i)
# Reset the device iterator resources with the new dataset.
ds_variant = ds._variant_tensor
gen_dataset_ops.make_iterator(
ds_variant, self._device_iterators[i]._iterator_resource)
def make_initializer(self, dataset, name=None):
"""Returns a `tf.Operation` that initializes this iterator on `dataset`.
Args:
dataset: A `Dataset` with compatible structure to 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 structure.
"""
with ops.name_scope(name, "make_initializer") as name:
nest.assert_same_structure(self._output_types, dataset.output_types)
nest.assert_same_structure(self._output_shapes, dataset.output_shapes)
for iterator_dtype, dataset_dtype in zip(
nest.flatten(self._output_types), nest.flatten(dataset.output_types)):
if iterator_dtype != dataset_dtype:
raise TypeError(
"Expected output types %r but got dataset with output types %r." %
(self._output_types, dataset.output_types))
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("Expected output shapes compatible with %r but got "
"dataset with output shapes %r." %
(self._output_shapes, dataset.output_shapes))
with ops.colocate_with(self._iterator_resource):
return gen_dataset_ops.make_iterator(
dataset._as_variant_tensor(), self._iterator_resource, name=name) # pylint: disable=protected-access
def make_initializer(self, dataset, name=None):
"""Returns a `tf.Operation` that initializes this iterator on `dataset`.
Args:
dataset: A `Dataset` with compatible structure to 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 structure.
"""
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(
"Expected output classes %r but got dataset with output class %r."
% (self.output_classes, dataset_output_classes))
for iterator_dtype, dataset_dtype in zip(
nest.flatten(self.output_types),
nest.flatten(dataset_output_types)):
if iterator_dtype != dataset_dtype:
raise TypeError(
"Expected output types %r but got dataset with output types %r."
% (self.output_types, dataset_output_types))
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(
"Expected output shapes compatible with %r but got "
"dataset with output shapes %r." %
(self.output_shapes, dataset_output_shapes))
with ops.colocate_with(dataset._variant_tensor):
return gen_dataset_ops.make_iterator(dataset._variant_tensor,
self._iterator_resource,
name=name) # pylint: disable=protected-access
def _eager_reset(self):
"""Resets the MultiDeviceIterator in eager mode."""
if not context.executing_eagerly():
raise ValueError("Eager reset is only supported in eager mode.")
# pylint: disable=protected-access
self._incarnation_id = gen_dataset_ops.multi_device_iterator_init(
self._dataset._variant_tensor,
self._multi_device_iterator_resource,
max_buffer_size=self._max_buffer_size)
for i, device in enumerate(self._devices):
with ops.device(device):
ds = self._create_device_dataset(i)
# Reset the device iterator resources with the new dataset.
ds_variant = ds._variant_tensor
gen_dataset_ops.make_iterator(
ds_variant, self._device_iterators[i]._iterator_resource)
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.executing_eagerly():
raise RuntimeError(
"{} objects can only be used when eager execution is enabled, use "
"tf.data.Dataset.make_initializable_iterator or "
"tf.data.Dataset.make_one_shot_iterator for graph construction"
.format(type(self)))
with ops.device("/device:CPU:0"):
ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access
self._output_classes = dataset.output_classes
self._output_types = dataset.output_types
self._output_shapes = dataset.output_shapes
self._flat_output_types = nest.flatten(
sparse.as_dense_types(self._output_types,
self._output_classes))
self._flat_output_shapes = nest.flatten(
sparse.as_dense_shapes(self._output_shapes,
self._output_classes))
self._resource = gen_dataset_ops.iterator(
shared_name="",
container=_generate_shared_name("eageriterator"),
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
# Delete the resource when this object is deleted
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._resource, handle_device="/device:CPU:0")
self._device = context.context().device_name
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.executing_eagerly():
raise RuntimeError(
"{} objects can only be used when eager execution is enabled, use "
"tf.data.Dataset.make_initializable_iterator or "
"tf.data.Dataset.make_one_shot_iterator for graph construction".
format(type(self)))
with ops.device("/device:CPU:0"):
ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access
self._output_classes = dataset.output_classes
self._output_types = dataset.output_types
self._output_shapes = dataset.output_shapes
self._flat_output_types = nest.flatten(
sparse.as_dense_types(self._output_types, self._output_classes))
self._flat_output_shapes = nest.flatten(
sparse.as_dense_shapes(self._output_shapes, self._output_classes))
self._resource = gen_dataset_ops.iterator(
shared_name="",
container=_generate_shared_name("eageriterator"),
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
# Delete the resource when this object is deleted
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._resource, handle_device="/device:CPU:0")
self._device = context.context().device_name
def _create_iterator(self, dataset):
# pylint: disable=protected-access
dataset = dataset._apply_options()
ds_variant = dataset._variant_tensor
self._structure = dataset._element_structure
self._flat_output_types = self._structure._flat_types
self._flat_output_shapes = self._structure._flat_shapes
with ops.colocate_with(ds_variant):
self._iterator_resource, self._deleter = (
gen_dataset_ops.anonymous_iterator_v2(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes))
gen_dataset_ops.make_iterator(ds_variant, self._iterator_resource)
# Delete the resource when this object is deleted
self._resource_deleter = IteratorResourceDeleter(
handle=self._iterator_resource,
device=self._device,
deleter=self._deleter)
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.executing_eagerly():
raise RuntimeError(
"{} objects can only be used when eager execution is enabled, use "
"tf.data.Dataset.make_initializable_iterator or "
"tf.data.Dataset.make_one_shot_iterator for graph construction"
.format(type(self)))
self._device = context.context().device_name
with ops.device("/cpu:0"):
# pylint: disable=protected-access
dataset = dataset._apply_options()
ds_variant = dataset._variant_tensor
self._structure = structure_lib.convert_legacy_structure(
dataset.output_types, dataset.output_shapes,
dataset.output_classes)
self._flat_output_types = self._structure._flat_types
self._flat_output_shapes = self._structure._flat_shapes
with ops.colocate_with(ds_variant):
self._resource = gen_dataset_ops.anonymous_iterator(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
# Delete the resource when this object is deleted
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._resource, handle_device=self._device)
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.executing_eagerly():
raise RuntimeError(
"{} objects can only be used when eager execution is enabled, use "
"tf.data.Dataset.make_initializable_iterator or "
"tf.data.Dataset.make_one_shot_iterator for graph construction".
format(type(self)))
self._device = context.context().device_name
with ops.device("/cpu:0"):
# pylint: disable=protected-access
dataset = dataset._apply_options()
ds_variant = dataset._variant_tensor
self._structure = structure_lib.convert_legacy_structure(
dataset.output_types, dataset.output_shapes, dataset.output_classes)
self._flat_output_types = self._structure._flat_types
self._flat_output_shapes = self._structure._flat_shapes
with ops.colocate_with(ds_variant):
self._resource = gen_dataset_ops.anonymous_iterator(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
# Delete the resource when this object is deleted
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._resource, handle_device=self._device)
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(
**dataset_ops.flat_structure(self._input_dataset))
with ops.control_dependencies(
[gen_dataset_ops.make_iterator(ds_variant, resource)]):
return gen_dataset_ops.iterator_to_string_handle(resource)
def _eager_reset(self):
"""Resets the MultiDeviceIterator in eager mode."""
if not ops.executing_eagerly_outside_functions():
raise ValueError(
"Resetting a multi-device iterator is only supported in the eager "
"mode.")
# pylint: disable=protected-access
self._incarnation_id = gen_dataset_ops.multi_device_iterator_init(
self._dataset._variant_tensor,
self._multi_device_iterator_resource,
max_buffer_size=self._max_buffer_size)
for i, device in enumerate(self._devices):
with ops.device(device):
ds = _create_device_dataset(self._prototype_device_datasets[i],
self._incarnation_id,
self._prefetch_buffer_size,
self._experimental_slack)
# Reset the device iterator resources with the new dataset.
ds_variant = ds._variant_tensor
gen_dataset_ops.make_iterator(
ds_variant, self._device_iterators[i]._iterator_resource)
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)
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
self._device = context.context().device_name
with ops.device("/cpu:0"):
# pylint: disable=protected-access
dataset = dataset._apply_options()
ds_variant = dataset._variant_tensor
self._structure = dataset._element_structure
self._flat_output_types = self._structure._flat_types
self._flat_output_shapes = self._structure._flat_shapes
with ops.colocate_with(ds_variant):
self._iterator_resource, self._deleter = (
gen_dataset_ops.anonymous_iterator_v2(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes))
gen_dataset_ops.make_iterator(ds_variant,
self._iterator_resource)
# Delete the resource when this object is deleted
self._resource_deleter = IteratorResourceDeleter(
handle=self._iterator_resource,
device=self._device,
deleter=self._deleter)
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
self._device = context.context().device_name
with ops.device("/cpu:0"):
# pylint: disable=protected-access
dataset = dataset._apply_options()
ds_variant = dataset._variant_tensor
self._structure = dataset._element_structure
self._flat_output_types = self._structure._flat_types
self._flat_output_shapes = self._structure._flat_shapes
with ops.colocate_with(ds_variant):
self._iterator_resource, self._deleter = (
gen_dataset_ops.anonymous_iterator_v2(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes))
gen_dataset_ops.make_iterator(ds_variant, self._iterator_resource)
# Delete the resource when this object is deleted
self._resource_deleter = IteratorResourceDeleter(
handle=self._iterator_resource,
device=self._device,
deleter=self._deleter)
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(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
with ops.control_dependencies(
[gen_dataset_ops.make_iterator(ds_variant, resource)]):
return gen_dataset_ops.iterator_to_string_handle(resource)
def _init_func():
"""Creates an iterator for the input dataset.
Returns:
A `string` tensor that encapsulates the iterator created.
"""
# pylint: disable=protected-access
ds_variant = self._input_dataset._as_variant_tensor()
resource = gen_dataset_ops.anonymous_iterator(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
with ops.control_dependencies(
[gen_dataset_ops.make_iterator(ds_variant, resource)]):
return gen_dataset_ops.iterator_to_string_handle(resource)
def _init_func():
"""Creates an iterator for the input dataset.
Returns:
A `string` tensor that encapsulates the iterator created.
"""
# pylint: disable=protected-access
ds_variant = self._input_dataset._as_variant_tensor()
resource = gen_dataset_ops.anonymous_iterator(
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
with ops.control_dependencies(
[gen_dataset_ops.make_iterator(ds_variant, resource)]):
return gen_dataset_ops.iterator_to_string_handle(resource)
def make_initializer(self, dataset, name=None):
"""Returns a `tf.Operation` that initializes this iterator on `dataset`.
Args:
dataset: A `Dataset` with compatible structure to 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 structure.
"""
with ops.name_scope(name, "make_initializer") as name:
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(
"Expected output classes %r but got dataset with output class %r."
% (self._output_classes, dataset.output_classes))
for iterator_dtype, dataset_dtype in zip(
nest.flatten(self._output_types),
nest.flatten(dataset.output_types)):
if iterator_dtype != dataset_dtype:
raise TypeError(
"Expected output types %r but got dataset with output types %r."
% (self._output_types, dataset.output_types))
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(
"Expected output shapes compatible with %r but got "
"dataset with output shapes %r." %
(self._output_shapes, dataset.output_shapes))
with ops.colocate_with(self._iterator_resource):
return gen_dataset_ops.make_iterator(dataset._as_variant_tensor(),
self._iterator_resource,
name=name) # pylint: disable=protected-access
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.in_eager_mode():
raise RuntimeError(
"{} objects can only be used when eager execution is enabled, use "
"tf.data.Dataset.make_iterator or "
"tf.data.Dataset.make_one_shot_iterator for graph construction"
.format(type(self)))
with ops.device("/device:CPU:0"):
ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access
self._output_types = dataset.output_types
self._output_shapes = dataset.output_shapes
self._flat_output_types = nest.flatten(dataset.output_types)
self._flat_output_shapes = nest.flatten(dataset.output_shapes)
self._resource = gen_dataset_ops.iterator(
container="",
shared_name=_generate_shared_name("eager_iterator"),
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
# Delete the resource when this object is deleted
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._resource, handle_device="/device:CPU:0")
self._device = context.context().device_name
self._buffer_resource_handle = None
if not context.context().device_spec.device_type:
is_remote_device = False
else:
is_remote_device = context.context(
).device_spec.device_type != "CPU"
if is_remote_device:
with ops.device("/device:CPU:0"):
iter_string_handle = gen_dataset_ops.iterator_to_string_handle(
self._resource)
@function.Defun(dtypes.string)
def remote_fn(h):
remote_iterator = iterator_ops.Iterator.from_string_handle(
h, self._output_types, self._output_shapes)
return remote_iterator.get_next()
remote_fn.add_to_graph(None)
target = constant_op.constant("/device:CPU:0")
with ops.device(self._device):
self._buffer_resource_handle = prefetching_ops.function_buffering_resource(
string_arg=iter_string_handle,
f=remote_fn,
target_device=target,
buffer_size=10,
thread_pool_size=1,
container="",
shared_name=_generate_shared_name(
"function_buffer_resource"))
self._buffer_resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._buffer_resource_handle,
handle_device=self._device)
def __init__(self, dataset):
"""Creates a new iterator over the given dataset.
For example:
```python
dataset = tf.data.Dataset.range(4)
for x in Iterator(dataset):
print(x)
```
Tensors produced will be placed on the device on which this iterator object
was created.
Args:
dataset: A `tf.data.Dataset` object.
Raises:
RuntimeError: When invoked without eager execution enabled.
"""
if not context.in_eager_mode():
raise RuntimeError(
"{} objects can only be used when eager execution is enabled, use "
"tf.data.Dataset.make_iterator or "
"tf.data.Dataset.make_one_shot_iterator for graph construction".
format(type(self)))
with ops.device("/device:CPU:0"):
ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access
self._output_types = dataset.output_types
self._output_shapes = dataset.output_shapes
self._flat_output_types = nest.flatten(dataset.output_types)
self._flat_output_shapes = nest.flatten(dataset.output_shapes)
self._resource = gen_dataset_ops.iterator(
container="",
shared_name=_generate_shared_name("eager_iterator"),
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
gen_dataset_ops.make_iterator(ds_variant, self._resource)
# Delete the resource when this object is deleted
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._resource, handle_device="/device:CPU:0")
self._device = context.context().device_name
self._buffer_resource_handle = None
if not context.context().device_spec.device_type:
is_remote_device = False
else:
is_remote_device = context.context().device_spec.device_type != "CPU"
if is_remote_device:
with ops.device("/device:CPU:0"):
iter_string_handle = gen_dataset_ops.iterator_to_string_handle(
self._resource)
@function.Defun(dtypes.string)
def remote_fn(h):
remote_iterator = iterator_ops.Iterator.from_string_handle(
h, self._output_types, self._output_shapes)
return remote_iterator.get_next()
remote_fn.add_to_graph(None)
target = constant_op.constant("/device:CPU:0")
with ops.device(self._device):
self._buffer_resource_handle = prefetching_ops.function_buffering_resource(
string_arg=iter_string_handle,
f=remote_fn,
target_device=target,
buffer_size=10,
thread_pool_size=1,
container="",
shared_name=_generate_shared_name("function_buffer_resource"))
self._buffer_resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._buffer_resource_handle, handle_device=self._device)
