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 _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. """ 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 create_iter(): return gen_dataset_ops.anonymous_iterator_v2( output_types=[dtypes.float32], output_shapes=[[]])