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): if use_anonymous_iterator_v3(): 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) else: 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, deleter=self._deleter)
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 create_iter(): with ops.device("/device:TPU:0"): return gen_dataset_ops.anonymous_iterator_v3( output_types=[dtypes.float32], output_shapes=[[]])