def from_value(value):
"""Returns an `Optional` that wraps the given value.
Args:
value: A nested structure of `tf.Tensor` and/or `tf.SparseTensor` objects.
Returns:
An `Optional` that wraps `value`.
"""
# TODO(b/110122868): Consolidate this destructuring logic with the
# similar code in `Dataset.from_tensors()`.
with ops.name_scope("optional") as scope:
with ops.name_scope("value"):
value = nest.pack_sequence_as(value, [
sparse_tensor_lib.SparseTensor.from_value(t)
if sparse_tensor_lib.is_sparse(t) else ops.convert_to_tensor(
t, name="component_%d" % i)
for i, t in enumerate(nest.flatten(value))
])
encoded_value = nest.flatten(sparse.serialize_sparse_tensors(value))
output_classes = sparse.get_classes(value)
output_shapes = nest.pack_sequence_as(
value, [t.get_shape() for t in nest.flatten(value)])
output_types = nest.pack_sequence_as(
value, [t.dtype for t in nest.flatten(value)])
return _OptionalImpl(
gen_dataset_ops.optional_from_value(encoded_value, name=scope),
output_shapes, output_types, output_classes)
def tf_finalize_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
for arg, shape in zip(
args,
nest.flatten(
sparse.as_dense_shapes(self._state_shapes, self._state_classes))):
arg.set_shape(shape)
nested_args = nest.pack_sequence_as(self._state_types, args)
nested_args = sparse.deserialize_sparse_tensors(
nested_args, self._state_types, self._state_shapes,
self._state_classes)
ret = finalize_func(nested_args)
# Convert any `SparseTensorValue`s to `SparseTensor`s and all other
# values to tensors.
ret = nest.pack_sequence_as(ret, [
sparse_tensor.SparseTensor.from_value(t)
if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
for t in nest.flatten(ret)
])
self._output_classes = sparse.get_classes(ret)
self._output_shapes = nest.pack_sequence_as(
ret, [t.get_shape() for t in nest.flatten(ret)])
self._output_types = nest.pack_sequence_as(
ret, [t.dtype for t in nest.flatten(ret)])
dataset_ops._warn_if_collections("tf.contrib.data.group_by_reducer()") # pylint: disable=protected-access
# Serialize any sparse tensors.
ret = nest.pack_sequence_as(
ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))])
return nest.flatten(ret)
def testFlattenAndPack(self):
structure = ((3, 4), 5, (6, 7, (9, 10), 8))
flat = ["a", "b", "c", "d", "e", "f", "g", "h"]
self.assertEqual(nest.flatten(structure), [3, 4, 5, 6, 7, 9, 10, 8])
self.assertEqual(
nest.pack_sequence_as(structure, flat), (("a", "b"), "c",
("d", "e", ("f", "g"), "h")))
point = collections.namedtuple("Point", ["x", "y"])
structure = (point(x=4, y=2), ((point(x=1, y=0),),))
flat = [4, 2, 1, 0]
self.assertEqual(nest.flatten(structure), flat)
restructured_from_flat = nest.pack_sequence_as(structure, flat)
self.assertEqual(restructured_from_flat, structure)
self.assertEqual(restructured_from_flat[0].x, 4)
self.assertEqual(restructured_from_flat[0].y, 2)
self.assertEqual(restructured_from_flat[1][0][0].x, 1)
self.assertEqual(restructured_from_flat[1][0][0].y, 0)
self.assertEqual([5], nest.flatten(5))
self.assertEqual([np.array([5])], nest.flatten(np.array([5])))
self.assertEqual("a", nest.pack_sequence_as(5, ["a"]))
self.assertEqual(
np.array([5]), nest.pack_sequence_as("scalar", [np.array([5])]))
with self.assertRaisesRegexp(ValueError, "Structure is a scalar"):
nest.pack_sequence_as("scalar", [4, 5])
with self.assertRaisesRegexp(TypeError, "flat_sequence"):
nest.pack_sequence_as([4, 5], "bad_sequence")
with self.assertRaises(ValueError):
nest.pack_sequence_as([5, 6, [7, 8]], ["a", "b", "c"])
def tf_finalize_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
for arg, shape in zip(
args,
nest.flatten(
sparse.as_dense_shapes(self._state_shapes, self._state_classes))):
arg.set_shape(shape)
nested_args = nest.pack_sequence_as(self._state_types, args)
nested_args = sparse.deserialize_sparse_tensors(
nested_args, self._state_types, self._state_shapes,
self._state_classes)
ret = finalize_func(nested_args)
# Convert any `SparseTensorValue`s to `SparseTensor`s and all other
# values to tensors.
ret = nest.pack_sequence_as(ret, [
sparse_tensor.SparseTensor.from_value(t)
if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
for t in nest.flatten(ret)
])
self._output_classes = sparse.get_classes(ret)
self._output_shapes = nest.pack_sequence_as(
ret, [t.get_shape() for t in nest.flatten(ret)])
self._output_types = nest.pack_sequence_as(
ret, [t.dtype for t in nest.flatten(ret)])
# Serialize any sparse tensors.
ret = nest.pack_sequence_as(
ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))])
return nest.flatten(ret)
def tf_reduce_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
for arg, shape in zip(
args,
nest.flatten(
sparse.as_dense_shapes(self._state_shapes, self._state_classes))
+ nest.flatten(
sparse.as_dense_shapes(input_dataset.output_shapes,
input_dataset.output_classes))):
arg.set_shape(shape)
pivot = len(nest.flatten(self._state_shapes))
nested_state_args = nest.pack_sequence_as(self._state_types,
args[:pivot])
nested_state_args = sparse.deserialize_sparse_tensors(
nested_state_args, self._state_types, self._state_shapes,
self._state_classes)
nested_input_args = nest.pack_sequence_as(input_dataset.output_types,
args[pivot:])
nested_input_args = sparse.deserialize_sparse_tensors(
nested_input_args, input_dataset.output_types,
input_dataset.output_shapes, input_dataset.output_classes)
ret = reduce_func(nested_state_args, nested_input_args)
# Convert any `SparseTensorValue`s to `SparseTensor`s and all other
# values to tensors.
ret = nest.pack_sequence_as(ret, [
sparse_tensor.SparseTensor.from_value(t)
if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
for t in nest.flatten(ret)
])
# Extract shape information from the returned values.
flat_new_state = nest.flatten(ret)
flat_new_state_shapes.extend([t.get_shape() for t in flat_new_state])
# Extract and validate type information from the returned values.
for t, dtype in zip(flat_new_state, nest.flatten(self._state_types)):
if t.dtype != dtype:
raise TypeError(
"The element types for the new state must match the initial "
"state. Expected %s; got %s." %
(self._state_types,
nest.pack_sequence_as(self._state_types,
[t.dtype for t in flat_new_state])))
dataset_ops._warn_if_collections("tf.contrib.data.group_by_reducer()") # pylint: disable=protected-access
# Serialize any sparse tensors.
ret = nest.pack_sequence_as(
ret,
[t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))])
return nest.flatten(ret)
def testPackDictOrder(self):
"""Packing orders dicts by key, including OrderedDicts."""
ordered = collections.OrderedDict([("d", 0), ("b", 0), ("a", 0), ("c", 0)])
plain = {"d": 0, "b": 0, "a": 0, "c": 0}
seq = [0, 1, 2, 3]
ordered_reconstruction = nest.pack_sequence_as(ordered, seq)
plain_reconstruction = nest.pack_sequence_as(plain, seq)
self.assertEqual(
collections.OrderedDict([("d", 3), ("b", 1), ("a", 0), ("c", 2)]),
ordered_reconstruction)
self.assertEqual({"d": 3, "b": 1, "a": 0, "c": 2}, plain_reconstruction)
def tf_key_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
# Pass in shape information from the input_dataset.
dense_shapes = sparse.as_dense_shapes(input_dataset.output_shapes,
input_dataset.output_classes)
for arg, shape in zip(args, nest.flatten(dense_shapes)):
arg.set_shape(shape)
nested_args = nest.pack_sequence_as(input_dataset.output_types, args)
nested_args = sparse.deserialize_sparse_tensors(
nested_args, input_dataset.output_types, input_dataset.output_shapes,
input_dataset.output_classes)
# pylint: disable=protected-access
if dataset_ops._should_unpack_args(nested_args):
ret = key_func(*nested_args)
# pylint: enable=protected-access
else:
ret = key_func(nested_args)
ret = ops.convert_to_tensor(ret)
if ret.dtype != dtypes.int64 or ret.get_shape() != tensor_shape.scalar():
raise ValueError(
"`key_func` must return a single tf.int64 tensor. "
"Got type=%s and shape=%s" % (ret.dtype, ret.get_shape()))
dataset_ops._warn_if_collections("tf.contrib.data.group_by_reducer()") # pylint: disable=protected-access
return ret
def normalize_tensors(tensors):
"""Converts a nested structure of tensor-like objects to tensors.
* `SparseTensor`-like inputs are converted to `SparseTensor`.
* `TensorArray` inputs are passed through.
* Everything else is converted to a dense `Tensor`.
Args:
tensors: A nested structure of tensor-like, list,
`SparseTensor`, `SparseTensorValue`, or `TensorArray` objects.
Returns:
A nested structure of tensor, `SparseTensor`, or `TensorArray` objects.
"""
flat_tensors = nest.flatten(tensors)
prepared = []
with ops.name_scope("normalize_tensors"):
for i, t in enumerate(flat_tensors):
if sparse_tensor_lib.is_sparse(t):
prepared.append(sparse_tensor_lib.SparseTensor.from_value(t))
elif ragged_tensor.is_ragged(t):
prepared.append(
ragged_tensor.convert_to_tensor_or_ragged_tensor(
t, name="component_%d" % i))
elif isinstance(t, tensor_array_ops.TensorArray):
prepared.append(t)
else:
prepared.append(ops.convert_to_tensor(t, name="component_%d" % i))
return nest.pack_sequence_as(tensors, prepared)
def _check_shape(*elements):
flatten_tensors = nest.flatten(elements)
flatten_shapes = nest.flatten(expected_shapes)
checked_tensors = [with_shape(shape, tensor)
for shape, tensor in zip(flatten_shapes,
flatten_tensors)]
return nest.pack_sequence_as(elements, checked_tensors)
def _next_internal(self):
"""Returns a nested structure of `tf.Tensor`s containing the next element.
"""
with ops.device(self._device):
if self._buffer_resource_handle is not None:
ret = prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=self._buffer_resource_handle,
output_types=self._flat_output_types)
else:
# TODO(ashankar): Consider removing this ops.device() contextmanager
# and instead mimic ops placement in graphs: Operations on resource
# handles execute on the same device as where the resource is placed.
# NOTE(mrry): Here we use the "_sync" variant of `iterator_get_next`
# because in eager mode this code will run synchronously on the calling
# thread. Therefore we do not need to make a defensive context switch
# to a background thread, and can achieve a small constant performance
# boost by invoking the iterator synchronously.
ret = gen_dataset_ops.iterator_get_next_sync(
self._resource,
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
return sparse.deserialize_sparse_tensors(
nest.pack_sequence_as(self._output_types, ret), self._output_types,
self._output_shapes, self._output_classes)
def get_next(self, name=None):
"""See `tf.data.Iterator.get_next`."""
self._get_next_call_count += 1
if self._get_next_call_count > iterator_ops.GET_NEXT_CALL_WARNING_THRESHOLD:
warnings.warn(iterator_ops.GET_NEXT_CALL_WARNING_MESSAGE)
flat_result = []
# TODO(priyag): This will fail if the input size (typically number of
# batches) is not divisible by number of devices.
# How do we handle that more gracefully / let the user know?
for buffer_resource in self._buffering_resources:
flat_ret = gen_dataset_ops.function_buffering_resource_get_next(
buffer_resource,
output_types=data_nest.flatten(sparse.as_dense_types(
self.output_types, self.output_classes)), name=name)
ret = sparse.deserialize_sparse_tensors(
data_nest.pack_sequence_as(self.output_types, flat_ret),
self.output_types, self.output_shapes, self.output_classes)
for tensor, shape in zip(
data_nest.flatten(ret), data_nest.flatten(self.output_shapes)):
if isinstance(tensor, ops.Tensor):
tensor.set_shape(shape)
flat_result.append(ret)
return nest.pack_sequence_as(self._devices, flat_result)
def generator_map_fn(iterator_id_t):
"""Generates the next element from iterator with ID `iterator_id_t`.
We map this function across an infinite repetition of the
`iterator_id_t`, and raise `StopIteration` to terminate the iteration.
Args:
iterator_id_t: A `tf.int64` tensor whose value uniquely identifies
the iterator in `generator_state` from which to generate an element.
Returns:
A nested structure of tensors representing an element from the iterator.
"""
def generator_py_func(iterator_id):
"""A `py_func` that will be called to invoke the iterator."""
try:
values = next(generator_state.get_iterator(iterator_id))
except StopIteration:
generator_state.iterator_completed(iterator_id)
raise StopIteration("Iteration finished.")
# Use the same _convert function from the py_func() implementation to
# convert the returned values to arrays early, so that we can inspect
# their values.
# pylint: disable=protected-access
ret_arrays = [
script_ops.FuncRegistry._convert(ret, dtype=dtype.as_numpy_dtype)
for ret, dtype in zip(nest.flatten_up_to(output_types, values),
flattened_types)
]
# pylint: enable=protected-access
# Additional type and shape checking to ensure that the components
# of the generated element match the `output_types` and `output_shapes`
# arguments.
for (ret_array, expected_dtype, expected_shape) in zip(
ret_arrays, flattened_types, flattened_shapes):
if ret_array.dtype != expected_dtype.as_numpy_dtype:
raise TypeError(
"`generator` yielded an element of type %s where an element "
"of type %s was expected." % (ret_array.dtype,
expected_dtype.as_numpy_dtype))
if not expected_shape.is_compatible_with(ret_array.shape):
raise ValueError(
"`generator` yielded an element of shape %s where an element "
"of shape %s was expected." % (ret_array.shape, expected_shape))
return ret_arrays
flat_values = script_ops.py_func(
generator_py_func, [iterator_id_t], flattened_types, stateful=True)
# The `py_func()` op drops the inferred shapes, so we add them back in
# here.
if output_shapes is not None:
for ret_t, shape in zip(flat_values, flattened_shapes):
ret_t.set_shape(shape)
return nest.pack_sequence_as(output_types, flat_values)
def tf_map_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
# Pass in shape information from the input_dataset.
dense_shapes = sparse.as_dense_shapes(input_dataset.output_shapes,
input_dataset.output_classes)
for arg, shape in zip(args, nest.flatten(dense_shapes)):
arg.set_shape(shape)
nested_args = nest.pack_sequence_as(input_dataset.output_types, args)
nested_args = sparse.deserialize_sparse_tensors(
nested_args, input_dataset.output_types, input_dataset.output_shapes,
input_dataset.output_classes)
if dataset_ops._should_unpack_args(nested_args): # pylint: disable=protected-access
dataset = map_func(*nested_args)
else:
dataset = map_func(nested_args)
if not isinstance(dataset, dataset_ops.Dataset):
raise TypeError("`map_func` must return a `Dataset` object.")
self._output_classes = dataset.output_classes
self._output_types = dataset.output_types
self._output_shapes = dataset.output_shapes
return dataset._as_variant_tensor() # pylint: disable=protected-access
def get_next(self, name=None):
"""Returns a nested structure of `tf.Tensor`s containing the next element.
Args:
name: (Optional.) A name for the created operation.
Returns:
A nested structure of `tf.Tensor` objects.
"""
self._get_next_call_count += 1
if self._get_next_call_count > GET_NEXT_CALL_WARNING_THRESHOLD:
warnings.warn(GET_NEXT_CALL_WARNING_MESSAGE)
return sparse.deserialize_sparse_tensors(
nest.pack_sequence_as(self._output_types,
gen_dataset_ops.iterator_get_next(
self._iterator_resource,
output_types=nest.flatten(
sparse.as_dense_types(
self._output_types,
self._output_classes)),
output_shapes=nest.flatten(
sparse.as_dense_shapes(
self._output_shapes,
self._output_classes)),
name=name)), self._output_types,
self._output_shapes, self._output_classes)
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_shapes`.
need_to_rerun = True
while need_to_rerun:
wrapped_func = dataset_ops.StructuredFunctionWrapper(
reduce_func,
self._transformation_name(),
input_classes=(self._state_classes, input_dataset.output_classes),
input_shapes=(self._state_shapes, input_dataset.output_shapes),
input_types=(self._state_types, input_dataset.output_types),
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(self._state_classes)):
if not issubclass(new_state_class, state_class):
raise TypeError(
"The element classes for the new state must match the initial "
"state. Expected %s; got %s." %
(self._state_classes, wrapped_func.output_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(self._state_types)):
if new_state_type != state_type:
raise TypeError(
"The element types for the new state must match the initial "
"state. Expected %s; got %s." %
(self._state_types, wrapped_func.output_types))
# Extract shape information from the returned values.
flat_state_shapes = nest.flatten(self._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:
self._state_shapes = nest.pack_sequence_as(self._state_shapes,
weakened_state_shapes)
self._reduce_func = wrapped_func.function
self._reduce_func.add_to_graph(ops.get_default_graph())
def _merge_output_shapes(original_shapes, expected_shapes):
flat_original_shapes = nest.flatten(original_shapes)
flat_new_shapes = nest.flatten_up_to(original_shapes, expected_shapes)
flat_merged_output_shapes = [
original_shape.merge_with(new_shape)
for original_shape, new_shape in zip(flat_original_shapes,
flat_new_shapes)]
return nest.pack_sequence_as(original_shapes, flat_merged_output_shapes)
def output_shapes(self):
ret = self._data_inputs[0].output_shapes
for data_input in self._data_inputs[1:]:
ret = nest.pack_sequence_as(ret, [
ts1.most_specific_compatible_shape(ts2) for (ts1, ts2) in zip(
nest.flatten(ret), nest.flatten(data_input.output_shapes))
])
return ret
def convert_legacy_structure(output_types, output_shapes, output_classes):
"""Returns a `Structure` that represents the given legacy structure.
This method provides a way to convert from the existing `Dataset` and
`Iterator` structure-related properties to a `Structure` object. A "legacy"
structure is represented by the `tf.data.Dataset.output_types`,
`tf.data.Dataset.output_shapes`, and `tf.data.Dataset.output_classes`
properties.
TODO(b/110122868): Remove this function once `Structure` is used throughout
`tf.data`.
Args:
output_types: A nested structure of `tf.DType` objects corresponding to
each component of a structured value.
output_shapes: A nested structure of `tf.TensorShape` objects
corresponding to each component a structured value.
output_classes: A nested structure of Python `type` objects corresponding
to each component of a structured value.
Returns:
A `Structure`.
Raises:
TypeError: If a structure cannot be built from the arguments, because one of
the component classes in `output_classes` is not supported.
"""
flat_types = nest.flatten(output_types)
flat_shapes = nest.flatten(output_shapes)
flat_classes = nest.flatten(output_classes)
flat_ret = []
for flat_type, flat_shape, flat_class in zip(flat_types, flat_shapes,
flat_classes):
if isinstance(flat_class, Structure):
flat_ret.append(flat_class)
elif issubclass(flat_class, sparse_tensor_lib.SparseTensor):
flat_ret.append(SparseTensorStructure(flat_type, flat_shape))
elif issubclass(flat_class, ops.Tensor):
flat_ret.append(TensorStructure(flat_type, flat_shape))
elif issubclass(flat_class, tensor_array_ops.TensorArray):
# We sneaked the dynamic_size and infer_shape into the legacy shape.
flat_ret.append(
TensorArrayStructure(
flat_type, flat_shape[2:],
dynamic_size=tensor_shape.dimension_value(flat_shape[0]),
infer_shape=tensor_shape.dimension_value(flat_shape[1])))
else:
# NOTE(mrry): Since legacy structures produced by iterators only
# comprise Tensors, SparseTensors, and nests, we do not need to
# support all structure types here.
raise TypeError(
"Could not build a structure for output class %r" % (flat_class,))
ret = nest.pack_sequence_as(output_classes, flat_ret)
if isinstance(ret, Structure):
return ret
else:
return NestedStructure(ret)
def next(self):
"""Return the next tf.Tensor from the dataset."""
try:
ret = gen_dataset_ops.iterator_get_next(
self._resource,
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
return nest.pack_sequence_as(self._output_types, ret)
except errors.OutOfRangeError:
raise StopIteration
def _tf_wrapper(idx):
flattened_types = nest.flatten(self.tf_dtypes)
flattened_shapes = nest.flatten(self.tf_shapes)
flat_values = tf.py_func(func=self._flatten_layer_op,
inp=[idx],
Tout=flattened_types)
for ret_t, shape in zip(flat_values, flattened_shapes):
# the actual returned numpy array shapes are not checked
ret_t.set_shape(shape)
return nest.pack_sequence_as(self.tf_dtypes, flat_values)
def make_padded_shapes(shapes, none_filler=None):
padded = []
for shape in nest.flatten(shapes):
shape = tensor_shape.TensorShape(shape)
shape = [
none_filler if d.value is None else d
for d in shape
]
padded.append(shape)
return nest.pack_sequence_as(shapes, padded)
def tf_scan_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
# Pass in shape information from the state and input_dataset.
for arg, shape in zip(
args,
flat_state_shapes + nest.flatten(input_dataset.output_shapes)):
arg.set_shape(shape)
pivot = len(flat_state_shapes)
old_state = nest.pack_sequence_as(self._initial_state, args[:pivot])
input_value = nest.pack_sequence_as(input_dataset.output_types,
args[pivot:])
ret = scan_func(old_state, input_value)
if not isinstance(ret, collections.Sequence) or len(ret) != 2:
raise TypeError("The scan function must return a pair comprising the "
"new state and the output value.")
new_state, output_value = ret
flat_new_state = [
ops.convert_to_tensor(t) for t in nest.flatten(new_state)
]
flat_output_value = [
ops.convert_to_tensor(t) for t in nest.flatten(output_value)
]
# Extract shape information from the returned values.
flat_new_state_shapes.extend([t.shape for t in flat_new_state])
self._output_shapes = nest.pack_sequence_as(
output_value, [t.shape for t in flat_output_value])
# Extract and validate type information from the returned values.
for t, dtype in zip(flat_new_state, flat_state_types):
if t.dtype != dtype:
raise TypeError(
"The element types for the new state must match the initial "
"state. Expected %s; got %s." %
(self._state_types, nest.pack_sequence_as(
self._state_types, [t.dtype for t in flat_new_state])))
self._output_types = nest.pack_sequence_as(
output_value, [t.dtype for t in flat_output_value])
return flat_new_state + flat_output_value
def _from_compatible_tensor_list(self, flat_value):
flat_ret = []
i = 0
for structure in self._flat_nested_structure:
num_flat_values = len(structure._flat_types)
sub_value = flat_value[i:i + num_flat_values]
flat_ret.append(structure._from_compatible_tensor_list(sub_value))
i += num_flat_values
return nest.pack_sequence_as(self._nested_structure, flat_ret)
def make_padded_shapes(shapes, none_filler=None):
padded = []
for shape in nest.flatten(shapes):
shape = tensor_shape.TensorShape(shape)
shape = [
none_filler if tensor_shape.dimension_value(d) is None else d
for d in shape
]
padded.append(shape)
return nest.pack_sequence_as(shapes, padded)
def next(self):
"""Return the next tf.Tensor from the dataset."""
try:
ret = gen_dataset_ops.iterator_get_next(
self._resource,
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
return nest.pack_sequence_as(self._output_types, ret)
except errors.OutOfRangeError:
raise StopIteration
def _tf_wrapper(idx):
flattened_types = nest.flatten(self.tf_dtypes)
flattened_shapes = nest.flatten(self.tf_shapes)
flat_values = tf.py_func(func=self._flatten_layer_op,
inp=[idx],
Tout=flattened_types)
for ret_t, shape in zip(flat_values, flattened_shapes):
# the actual returned numpy array shapes are not checked
ret_t.set_shape(shape)
return nest.pack_sequence_as(self.tf_dtypes, flat_values)
def _from_compatible_tensor_list(self, flat_value):
flat_ret = []
i = 0
for structure in self._flat_nested_structure:
num_flat_values = len(structure._flat_types)
sub_value = flat_value[i:i + num_flat_values]
flat_ret.append(structure._from_compatible_tensor_list(sub_value))
i += num_flat_values
return nest.pack_sequence_as(self._nested_structure, flat_ret)
def tf_scan_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
# Pass in shape information from the state and input_dataset.
for arg, shape in zip(
args,
flat_state_shapes + nest.flatten(input_dataset.output_shapes)):
arg.set_shape(shape)
pivot = len(flat_state_shapes)
old_state = nest.pack_sequence_as(self._initial_state, args[:pivot])
input_value = nest.pack_sequence_as(input_dataset.output_types,
args[pivot:])
ret = scan_func(old_state, input_value)
if not isinstance(ret, collections.Sequence) or len(ret) != 2:
raise TypeError("The scan function must return a pair comprising the "
"new state and the output value.")
new_state, output_value = ret
flat_new_state = [
ops.convert_to_tensor(t) for t in nest.flatten(new_state)
]
flat_output_value = [
ops.convert_to_tensor(t) for t in nest.flatten(output_value)
]
# Extract shape information from the returned values.
flat_new_state_shapes.extend([t.shape for t in flat_new_state])
self._output_shapes = nest.pack_sequence_as(
output_value, [t.shape for t in flat_output_value])
# Extract and validate type information from the returned values.
for t, dtype in zip(flat_new_state, flat_state_types):
if t.dtype != dtype:
raise TypeError(
"The element types for the new state must match the initial "
"state. Expected %s; got %s." %
(self._state_types, nest.pack_sequence_as(
self._state_types, [t.dtype for t in flat_new_state])))
self._output_types = nest.pack_sequence_as(
output_value, [t.dtype for t in flat_output_value])
return flat_new_state + flat_output_value
def _from_tensor_list(self, flat_value):
if len(flat_value) != len(self._flat_types):
raise ValueError("Expected %d flat values in NestedStructure but got %d."
% (len(self._flat_types), len(flat_value)))
flat_ret = []
for sub_value, structure in zip(flat_value, self._flat_nested_structure):
flat_ret.append(structure._from_tensor_list([sub_value]))
return nest.pack_sequence_as(self._nested_structure, flat_ret)
def __init__(self, dataset, output_types, output_shapes=None):
"""Creates a new dataset with the given output types and shapes.
The given `dataset` must have a structure that is convertible:
* `dataset.output_types` must be the same as `output_types` module nesting.
* Each shape in `dataset.output_shapes` must be compatible with each shape
in `output_shapes` (if given).
Note: This helper permits "unsafe casts" for shapes, equivalent to using
`tf.Tensor.set_shape()` where domain-specific knowledge is available.
Args:
dataset: A `Dataset` object.
output_types: A nested structure of `tf.DType` objects.
output_shapes: (Optional.) A nested structure of `tf.TensorShape` objects.
If omitted, the shapes will be inherited from `dataset`.
Raises:
ValueError: If either `output_types` or `output_shapes` is not compatible
with the structure of `dataset`.
"""
super(_RestructuredDataset, self).__init__()
self._dataset = dataset
# Validate that the types are compatible.
output_types = nest.map_structure(dtypes.as_dtype, output_types)
flat_original_types = nest.flatten(dataset.output_types)
flat_new_types = nest.flatten(output_types)
if flat_original_types != flat_new_types:
raise ValueError(
"Dataset with output types %r cannot be restructured to have output "
"types %r" % (dataset.output_types, output_types))
self._output_types = output_types
if output_shapes is None:
# Inherit shapes from the original `dataset`.
self._output_shapes = nest.pack_sequence_as(output_types,
nest.flatten(
dataset.output_shapes))
else:
# Validate that the shapes are compatible.
nest.assert_same_structure(output_types, output_shapes)
flat_original_shapes = nest.flatten(dataset.output_shapes)
flat_new_shapes = nest.flatten_up_to(output_types, output_shapes)
for original_shape, new_shape in zip(flat_original_shapes,
flat_new_shapes):
if not original_shape.is_compatible_with(new_shape):
raise ValueError(
"Dataset with output shapes %r cannot be restructured to have "
"incompatible output shapes %r" % (dataset.output_shapes,
output_shapes))
self._output_shapes = nest.map_structure_up_to(
output_types, tensor_shape.as_shape, output_shapes)
def normalize_element(element, dtypes=None):
"""Normalizes a nested structure of element components.
* Components matching `SparseTensorSpec` are converted to `SparseTensor`.
* Components matching `RaggedTensorSpec` are converted to `RaggedTensor`.
* Components matching `DatasetSpec` or `TensorArraySpec` are passed through.
* `CompositeTensor` components are passed through.
* All other components are converted to `Tensor`.
Args:
element: A nested structure of individual components.
dtypes: (Optional.) A nested structure of `tf.DType` objects corresponding
to each component of `element`. If specified, it will be used to set the
exact type of output tensor when converting input components which
are not tensors themselves (e.g. numpy arrays, native python types, etc.)
Returns:
A nested structure of `Tensor`, `Dataset`, `SparseTensor`, `RaggedTensor`,
or `TensorArray` objects.
"""
components = nest.flatten(element)
normalized_components = []
if dtypes is None:
flattened_dtypes = [None] * len(components)
else:
flattened_dtypes = nest.flatten(dtypes)
with ops.name_scope("normalize_element"):
# Imported here to avoid circular dependency.
from tensorflow.python.data.ops import dataset_ops # pylint: disable=g-import-not-at-top
for i, (t, dtype) in enumerate(zip(components, flattened_dtypes)):
try:
spec = type_spec_from_value(t, use_fallback=False)
except TypeError:
# TypeError indicates it was not possible to compute a `TypeSpec` for
# the value. As a fallback try converting the value to a tensor.
normalized_components.append(
ops.convert_to_tensor(t, name="component_%d" % i, dtype=dtype))
else:
if isinstance(spec, sparse_tensor.SparseTensorSpec):
normalized_components.append(sparse_tensor.SparseTensor.from_value(t))
elif isinstance(spec, ragged_tensor.RaggedTensorSpec):
normalized_components.append(
ragged_tensor.convert_to_tensor_or_ragged_tensor(
t, name="component_%d" % i))
elif isinstance(
spec, (tensor_array_ops.TensorArraySpec, dataset_ops.DatasetSpec)):
normalized_components.append(t)
elif isinstance(spec, NoneTensorSpec):
normalized_components.append(NoneTensor())
elif isinstance(t, composite_tensor.CompositeTensor):
normalized_components.append(t)
else:
normalized_components.append(
ops.convert_to_tensor(t, name="component_%d" % i, dtype=dtype))
return nest.pack_sequence_as(element, normalized_components)
def __init__(self, dataset, output_types, output_shapes=None):
"""Creates a new dataset with the given output types and shapes.
The given `dataset` must have a structure that is convertible:
* `dataset.output_types` must be the same as `output_types` module nesting.
* Each shape in `dataset.output_shapes` must be compatible with each shape
in `output_shapes` (if given).
Note: This helper permits "unsafe casts" for shapes, equivalent to using
`tf.Tensor.set_shape()` where domain-specific knowledge is available.
Args:
dataset: A `Dataset` object.
output_types: A nested structure of `tf.DType` objects.
output_shapes: (Optional.) A nested structure of `tf.TensorShape` objects.
If omitted, the shapes will be inherited from `dataset`.
Raises:
ValueError: If either `output_types` or `output_shapes` is not compatible
with the structure of `dataset`.
"""
super(_RestructuredDataset, self).__init__()
self._dataset = dataset
# Validate that the types are compatible.
output_types = nest.map_structure(dtypes.as_dtype, output_types)
flat_original_types = nest.flatten(dataset.output_types)
flat_new_types = nest.flatten(output_types)
if flat_original_types != flat_new_types:
raise ValueError(
"Dataset with output types %r cannot be restructured to have output "
"types %r" % (dataset.output_types, output_types))
self._output_types = output_types
if output_shapes is None:
# Inherit shapes from the original `dataset`.
self._output_shapes = nest.pack_sequence_as(output_types,
nest.flatten(
dataset.output_shapes))
else:
# Validate that the shapes are compatible.
nest.assert_same_structure(output_types, output_shapes)
flat_original_shapes = nest.flatten(dataset.output_shapes)
flat_new_shapes = nest.flatten_up_to(output_types, output_shapes)
for original_shape, new_shape in zip(flat_original_shapes,
flat_new_shapes):
if not original_shape.is_compatible_with(new_shape):
raise ValueError(
"Dataset with output shapes %r cannot be restructured to have "
"incompatible output shapes %r" % (dataset.output_shapes,
output_shapes))
self._output_shapes = nest.map_structure_up_to(
output_types, tensor_shape.as_shape, output_shapes)
def convert_legacy_structure(output_types, output_shapes, output_classes):
"""Returns a `Structure` that represents the given legacy structure.
This method provides a way to convert from the existing `Dataset` and
`Iterator` structure-related properties to a `Structure` object. A "legacy"
structure is represented by the `tf.data.Dataset.output_types`,
`tf.data.Dataset.output_shapes`, and `tf.data.Dataset.output_classes`
properties.
TODO(b/110122868): Remove this function once `Structure` is used throughout
`tf.data`.
Args:
output_types: A nested structure of `tf.DType` objects corresponding to
each component of a structured value.
output_shapes: A nested structure of `tf.TensorShape` objects
corresponding to each component a structured value.
output_classes: A nested structure of Python `type` objects corresponding
to each component of a structured value.
Returns:
A `Structure`.
Raises:
TypeError: If a structure cannot be built from the arguments, because one of
the component classes in `output_classes` is not supported.
"""
flat_types = nest.flatten(output_types)
flat_shapes = nest.flatten(output_shapes)
flat_classes = nest.flatten(output_classes)
flat_ret = []
for flat_type, flat_shape, flat_class in zip(flat_types, flat_shapes,
flat_classes):
if isinstance(flat_class, type_spec.TypeSpec):
flat_ret.append(flat_class)
elif issubclass(flat_class, sparse_tensor.SparseTensor):
flat_ret.append(SparseTensorStructure(flat_type, flat_shape))
elif issubclass(flat_class, ops.Tensor):
flat_ret.append(TensorStructure(flat_type, flat_shape))
elif issubclass(flat_class, tensor_array_ops.TensorArray):
# We sneaked the dynamic_size and infer_shape into the legacy shape.
flat_ret.append(
TensorArrayStructure(
flat_type,
flat_shape[2:],
dynamic_size=tensor_shape.dimension_value(flat_shape[0]),
infer_shape=tensor_shape.dimension_value(flat_shape[1])))
else:
# NOTE(mrry): Since legacy structures produced by iterators only
# comprise Tensors, SparseTensors, and nests, we do not need to
# support all structure types here.
raise TypeError("Could not build a structure for output class %r" %
(flat_class, ))
return nest.pack_sequence_as(output_classes, flat_ret)
def __init__(self,
selector_input,
data_inputs,
stop_on_empty_dataset=False):
self._selector_input = selector_input
self._data_inputs = list(data_inputs)
self._stop_on_empty_dataset = stop_on_empty_dataset
first_output_types = dataset_ops.get_legacy_output_types(
data_inputs[0])
first_output_classes = dataset_ops.get_legacy_output_classes(
data_inputs[0])
for i, data_input in enumerate(data_inputs[1:]):
if (dataset_ops.get_legacy_output_types(data_input) !=
first_output_types
or dataset_ops.get_legacy_output_classes(data_input) !=
first_output_classes):
raise TypeError(
"All datasets must have the same type and class.\n"
"dataset 0 vs dataset %s types: %s ; %s\n"
"classes: %s ; %s" %
(i + 1, first_output_types,
dataset_ops.get_legacy_output_types(data_input),
first_output_classes,
dataset_ops.get_legacy_output_classes(data_input)))
output_shapes = dataset_ops.get_legacy_output_shapes(
self._data_inputs[0])
for data_input in self._data_inputs[1:]:
output_shapes = nest.pack_sequence_as(output_shapes, [
ts1.most_specific_compatible_shape(ts2) for (ts1, ts2) in zip(
nest.flatten(output_shapes),
nest.flatten(
dataset_ops.get_legacy_output_shapes(data_input)))
])
self._element_spec = structure.convert_legacy_structure(
first_output_types, output_shapes, first_output_classes)
compat_kwargs = {}
if compat.forward_compatible(2021, 5,
14) or self._stop_on_empty_dataset:
compat_kwargs[
"stop_on_empty_dataset"] = self._stop_on_empty_dataset
# pylint: disable=protected-access
variant_tensor = (
gen_experimental_dataset_ops.directed_interleave_dataset(
self._selector_input._variant_tensor, [
data_input._variant_tensor
for data_input in self._data_inputs
], **compat_kwargs, **self._flat_structure))
super(_DirectedInterleaveDataset, self).__init__(variant_tensor)
def _from_tensor_list(self, flat_value):
if len(flat_value) != len(self._flat_types):
raise ValueError("Expected %d flat values in NestedStructure but got %d."
% (len(self._flat_types), len(flat_value)))
flat_ret = []
for sub_value, structure in zip(flat_value,
nest.flatten(self._nested_structure)):
flat_ret.append(structure._from_tensor_list([sub_value]))
return nest.pack_sequence_as(self._nested_structure, flat_ret)
def tpu_pre(loss):
"""Generate the TPU graph."""
del loss
values = self.infeed_queue[0].generate_dequeue_op(tpu_device=0)
unflattened_inputs = data_nest.pack_sequence_as(
self.feature_structure, values)
features = unflattened_inputs["features"]
labels = unflattened_inputs["labels"]
estimator_spec = model_fn(features, labels,
tf.estimator.ModeKeys.TRAIN, mparams)
return estimator_spec
def _prefetch_fn(handle):
"""Prefetches one element from `input_iterator`."""
remote_iterator = iterator_ops.Iterator.from_string_handle(
handle, self._input_iterator.output_types,
self._input_iterator.output_shapes,
self._input_iterator.output_classes)
ret = remote_iterator.get_next()
# Convert any `SparseTensorValue`s to `SparseTensor`s.
ret = nest.pack_sequence_as(ret, [
sparse_tensor_lib.SparseTensor.from_value(t)
if sparse_tensor_lib.is_sparse(t) else t for t in nest.flatten(ret)
])
# Serialize any sparse tensors and convert result to tensors.
ret = nest.pack_sequence_as(ret, [
ops.convert_to_tensor(t)
for t in nest.flatten(sparse.serialize_sparse_tensors(ret))
])
return nest.flatten(ret)
def tpu_train_step(loss):
"""Generate the TPU graph."""
del loss
values = self.infeed_queue[0].generate_dequeue_op(tpu_device=0)
unflattened_inputs = data_nest.pack_sequence_as(
self.feature_structure, values)
features = unflattened_inputs["features"]
labels = unflattened_inputs["labels"]
estimator_spec = model_fn(features, labels,
tf.estimator.ModeKeys.TRAIN, params)
loss, train_op = estimator_spec.loss, estimator_spec.train_op
with tf.control_dependencies([train_op]):
return tf.identity(loss)
def tf_init_func(key):
"""A wrapper for Defun that facilitates shape inference."""
key.set_shape([])
ret = init_func(key)
# Convert any `SparseTensorValue`s to `SparseTensor`s and all other
# values to tensors.
ret = nest.pack_sequence_as(ret, [
sparse_tensor.SparseTensor.from_value(t)
if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
for t in nest.flatten(ret)
])
self._state_classes = sparse.get_classes(ret)
self._state_shapes = nest.pack_sequence_as(
ret, [t.get_shape() for t in nest.flatten(ret)])
self._state_types = nest.pack_sequence_as(
ret, [t.dtype for t in nest.flatten(ret)])
# Serialize any sparse tensors.
ret = nest.pack_sequence_as(
ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))])
return nest.flatten(ret)
def _next_internal(self):
"""Returns a nested structure of `tf.Tensor`s containing the next element.
"""
if self._buffer_resource_handle is not None:
with ops.device(self._device):
ret = prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=self._buffer_resource_handle,
output_types=self._flat_output_types)
return sparse.deserialize_sparse_tensors(
nest.pack_sequence_as(self._output_types, ret),
self._output_types, self._output_shapes, self._output_classes)
else:
return super(Iterator, self)._next_internal()
def _next_internal(self):
"""Returns a nested structure of `tf.Tensor`s containing the next element.
"""
if self._buffer_resource_handle is not None:
with ops.device(self._device):
ret = prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=self._buffer_resource_handle,
output_types=self._flat_output_types)
return sparse.deserialize_sparse_tensors(
nest.pack_sequence_as(self._output_types, ret), self._output_types,
self._output_shapes, self._output_classes)
else:
return super(Iterator, self)._next_internal()
def grad(*dy):
"""Gradient function for build_rnn."""
dy = nest.pack_sequence_as(ret, dy)
def _continue(unused_theta, unused_dy, unused_dstate1,
unused_dtheta, unused_dinput, i):
return _should_continue(i, True)
dstate1, dtheta, dinput = tf.while_loop(_continue, _cell_grad_fn, [
theta,
dy,
{
k: tf.zeros_like(state0[k])
for k in state0 if k not in skipped_state
},
{
k: tf.zeros_like(theta[k])
for k in theta if k not in skipped_theta
},
{k: tf.zeros_like(inputs[k])
for k in inputs},
tf.zeros([], tf.int32) if reverse else max_length - 1,
])[2:5]
dtheta, dinput = _cell_grad_fn_with_state0(
state0, theta, dy, dstate1, dtheta, dinput, max_length -
1 if reverse else tf.zeros([], dtype=tf.int32))[3:5]
state0_h = tf.reshape(acc_state["h"],
[-1, theta["kernel"].shape[0]])
state0_atten = tf.reshape(acc_state["attention"], [
-1, theta["attention_kernel"].shape[0]
]) if "attention_kernel" in theta else None
grad = tf.reshape(dinput["rnn"], [-1, theta["kernel"].shape[1]])
if reverse:
state0_h = tf.split(state0_h, [batch_size, -1])[1]
grad = tf.split(grad, [-1, batch_size])[0]
else:
if state0_atten is not None:
state0_atten = tf.split(state0_atten, [-1, batch_size])[0]
state0_h = tf.split(state0_h, [-1, batch_size])[0]
grad = tf.split(grad, [batch_size, -1])[1]
if state0_atten is not None:
dtheta["attention_kernel"] = tf.matmul(
tf.transpose(state0_atten), grad)
dtheta["kernel"] = tf.matmul(tf.transpose(state0_h), grad)
if "memory_kernel" in orig_theta:
dtheta["memory_kernel"] = tf.zeros_like(
orig_theta["memory_kernel"])
dtheta["seq_mask"] = tf.zeros_like(orig_theta["seq_mask"])
return dinput, dtheta
def tf_finalize_func(*args):
"""A wrapper for Defun that facilitates shape inference."""
for arg, shape in zip(
args,
nest.flatten(
sparse.as_dense_shapes(self._state_shapes,
self._state_classes))):
arg.set_shape(shape)
nested_args = nest.pack_sequence_as(self._state_types, args)
nested_args = sparse.deserialize_sparse_tensors(
nested_args, self._state_types, self._state_shapes,
self._state_classes)
ret = finalize_func(nested_args)
# Convert any `SparseTensorValue`s to `SparseTensor`s and all other
# values to tensors.
ret = nest.pack_sequence_as(ret, [
sparse_tensor.SparseTensor.from_value(t)
if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
for t in nest.flatten(ret)
])
self._output_classes = sparse.get_classes(ret)
self._output_shapes = nest.pack_sequence_as(
ret, [t.get_shape() for t in nest.flatten(ret)])
self._output_types = nest.pack_sequence_as(
ret, [t.dtype for t in nest.flatten(ret)])
dataset_ops._warn_if_collections(
"tf.contrib.data.group_by_reducer()") # pylint: disable=protected-access
# Serialize any sparse tensors.
ret = nest.pack_sequence_as(ret, [
t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))
])
return nest.flatten(ret)
def next(self):
"""Return the next tf.Tensor from the dataset."""
try:
# TODO(ashankar): Consider removing this ops.device() contextmanager
# and instead mimic ops placement in graphs: Operations on resource
# handles execute on the same device as where the resource is placed.
with ops.device("/device:CPU:0"):
ret = gen_dataset_ops.iterator_get_next(
self._resource,
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
return nest.pack_sequence_as(self._output_types, ret)
except errors.OutOfRangeError:
raise StopIteration
def _testSaveRestoreUtility(self, start, break_range, stop):
path = self._iterator_checkpoint_prefix()
step = 0
meta_filename = path + "-%d.meta" % step
input_components = (np.tile(np.array([[1], [2], [3], [4]]), 20),
np.tile(np.array([[12], [13], [14], [15]]), 4))
to_concatenate_components = (np.tile(
np.array([[5], [6], [7], [8], [9]]),
20), np.tile(np.array([[16], [17], [18], [19], [20]]), 15))
with ops.Graph().as_default() as g:
init_op, get_next = self._build_graph(input_components,
to_concatenate_components)
saver = saver_lib.Saver()
with self.test_session(graph=g) as sess:
sess.run(init_op)
for i in range(start, break_range):
result = sess.run(get_next)
if i < 4:
for component, result_component in zip(
input_components, result):
self.assertAllEqual(component[i], result_component)
else:
for component, result_component in zip(
to_concatenate_components, result):
self.assertAllEqual(component[i - 4],
result_component)
saver.save(sess, path, step)
with ops.Graph().as_default() as g:
saver = saver_lib.import_meta_graph(meta_filename)
with self.test_session(graph=g) as sess:
get_next = nest.pack_sequence_as(
("a", "b"), ops.get_collection("get_next"))
saver.restore(sess,
saver_lib.latest_checkpoint(self.get_temp_dir()))
for i in range(break_range, stop):
result = sess.run(get_next)
if i < 4:
for component, result_component in zip(
input_components, result):
self.assertAllEqual(component[i], result_component)
else:
for component, result_component in zip(
to_concatenate_components, result):
self.assertAllEqual(component[i - 4],
result_component)
with self.assertRaises(errors.OutOfRangeError):
sess.run(get_next)
def _from_tensor_list(self, flat_value):
if len(flat_value) != len(self._flat_types):
raise ValueError("Expected %d flat values in NestedStructure but got %d."
% (len(self._flat_types), len(flat_value)))
flat_ret = []
i = 0
for structure in self._flat_nested_structure:
num_flat_values = len(structure._flat_types)
sub_value = flat_value[i:i + num_flat_values]
flat_ret.append(structure._from_tensor_list(sub_value))
i += num_flat_values
return nest.pack_sequence_as(self._nested_structure, flat_ret)
def _next_internal(self):
"""Returns a nested structure of `tf.Tensor`s containing the next element.
"""
# This runs in sync mode as iterators use an error status to communicate
# that there is no more data to iterate over.
# TODO(b/77291417): Fix
with context.execution_mode(context.SYNC):
with ops.device(self._device):
ret = ged_ops.experimental_function_buffering_resource_get_next(
function_buffer_resource=self._buffering_resource,
output_types=self._flat_output_types)
return sparse.deserialize_sparse_tensors(
nest.pack_sequence_as(self._output_types, ret), self._output_types,
self._output_shapes, self._output_classes)
def next(self):
"""Return the next tf.Tensor from the dataset."""
try:
# TODO(ashankar): Consider removing this ops.device() contextmanager
# and instead mimic ops placement in graphs: Operations on resource
# handles execute on the same device as where the resource is placed.
with ops.device("/device:CPU:0"):
ret = gen_dataset_ops.iterator_get_next(
self._resource,
output_types=self._flat_output_types,
output_shapes=self._flat_output_shapes)
return nest.pack_sequence_as(self._output_types, ret)
except errors.OutOfRangeError:
raise StopIteration
def __tensors_to_value(self, flat_value, from_tensor_list_fn):
if len(flat_value) != len(self._flat_tensor_specs):
raise ValueError(
"Expected %d flat values in NestedStructure but got %d." %
(len(self._flat_tensor_specs), len(flat_value)))
flat_ret = []
i = 0
for structure in self._flat_nested_structure:
num_flat_values = len(structure._flat_tensor_specs)
sub_value = flat_value[i:i + num_flat_values]
flat_ret.append(from_tensor_list_fn(structure, sub_value))
i += num_flat_values
return nest.pack_sequence_as(self._nested_structure, flat_ret)
def _next_internal(self):
"""Returns a nested structure of `tf.Tensor`s containing the next element.
"""
# This runs in sync mode as iterators use an error status to communicate
# that there is no more data to iterate over.
# TODO(b/77291417): Fix
with context.execution_mode(context.SYNC):
with ops.device(self._device):
ret = gen_dataset_ops.function_buffering_resource_get_next(
function_buffer_resource=self._buffering_resource,
output_types=self._flat_output_types)
return sparse.deserialize_sparse_tensors(
nest.pack_sequence_as(self._output_types, ret), self._output_types,
self._output_shapes, self._output_classes)
def tf_init_func(key):
"""A wrapper for Defun that facilitates shape inference."""
key.set_shape([])
ret = init_func(key)
# Convert any `SparseTensorValue`s to `SparseTensor`s and all other
# values to tensors.
ret = nest.pack_sequence_as(ret, [
sparse_tensor.SparseTensor.from_value(t)
if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
for t in nest.flatten(ret)
])
self._state_classes = sparse.get_classes(ret)
self._state_shapes = nest.pack_sequence_as(
ret, [t.get_shape() for t in nest.flatten(ret)])
self._state_types = nest.pack_sequence_as(
ret, [t.dtype for t in nest.flatten(ret)])
# Serialize any sparse tensors.
ret = nest.pack_sequence_as(ret, [
t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))
])
return nest.flatten(ret)
def _from_legacy_structure(output_types, output_shapes, output_classes):
"""Returns a `Structure` that represents the given legacy structure.
This method provides a way to convert from the existing `Dataset` and
`Iterator` structure-related properties to a `Structure` object.
TODO(b/110122868): Remove this method once `Structure` is used throughout
`tf.data`.
Args:
output_types: A nested structure of `tf.DType` objects corresponding to
each component of a structured value.
output_shapes: A nested structure of `tf.TensorShape` objects
corresponding to each component a structured value.
output_classes: A nested structure of Python `type` objects corresponding
to each component of a structured value.
Returns:
A `Structure`.
Raises:
TypeError: If a structure cannot be built the arguments, because one of
the component classes in `output_classes` is not supported.
"""
flat_types = nest.flatten(output_types)
flat_shapes = nest.flatten(output_shapes)
flat_classes = nest.flatten(output_classes)
flat_ret = []
for flat_type, flat_shape, flat_class in zip(flat_types, flat_shapes,
flat_classes):
if isinstance(flat_class, Structure):
flat_ret.append(flat_class)
elif issubclass(flat_class, sparse_tensor_lib.SparseTensor):
flat_ret.append(SparseTensorStructure(flat_type, flat_shape))
elif issubclass(flat_class, ops.Tensor):
flat_ret.append(TensorStructure(flat_type, flat_shape))
else:
# NOTE(mrry): Since legacy structures produced by iterators only
# comprise Tensors, SparseTensors, and nests, we do not need to
# support all structure types here.
raise TypeError(
"Could not build a structure for output class %r" %
flat_type)
ret = nest.pack_sequence_as(output_classes, flat_ret)
if isinstance(ret, Structure):
return ret
else:
return NestedStructure(ret)
def get_single_element(dataset):
"""Returns the single element in `dataset` as a nested structure of tensors.
This function enables you to use a @{tf.data.Dataset} in a stateless
"tensor-in tensor-out" expression, without creating a @{tf.data.Iterator}.
This can be useful when your preprocessing transformations are expressed
as a `Dataset`, and you want to use the transformation at serving time.
For example:
```python
input_batch = tf.placeholder(tf.string, shape=[BATCH_SIZE])
def preprocessing_fn(input_str):
# ...
return image, label
dataset = (tf.data.Dataset.from_tensor_slices(input_batch)
.map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
.batch(BATCH_SIZE))
image_batch, label_batch = tf.contrib.data.get_single_element(dataset)
```
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, dataset_ops.Dataset):
raise TypeError("`dataset` must be a `tf.data.Dataset` object.")
nested_ret = nest.pack_sequence_as(
dataset.output_types, gen_dataset_ops.dataset_to_single_element(
dataset._as_variant_tensor(), # pylint: disable=protected-access
output_types=nest.flatten(sparse.as_dense_types(
dataset.output_types, dataset.output_classes)),
output_shapes=nest.flatten(sparse.as_dense_shapes(
dataset.output_shapes, dataset.output_classes))))
return sparse.deserialize_sparse_tensors(
nested_ret, dataset.output_types, dataset.output_shapes,
dataset.output_classes)
def normalize_element(element):
"""Normalizes a nested structure of element components.
* Components matching `SparseTensorSpec` are converted to `SparseTensor`.
* Components matching `RaggedTensorSpec` are converted to `RaggedTensor`.
* Components matching `DatasetSpec` or `TensorArraySpec` are passed through.
* `CompositeTensor` components are passed through.
* All other components are converted to `Tensor`.
Args:
element: A nested structure of individual components.
Returns:
A nested structure of `Tensor`, `Dataset`, `SparseTensor`, `RaggedTensor`,
or `TensorArray` objects.
"""
components = nest.flatten(element)
normalized_components = []
with ops.name_scope("normalize_element"):
# Imported here to avoid circular dependency.
from tensorflow.python.data.ops import dataset_ops # pylint: disable=g-import-not-at-top
for i, t in enumerate(components):
try:
spec = type_spec_from_value(t, use_fallback=False)
except TypeError:
# TypeError indicates it was not possible to compute a `TypeSpec` for
# the value. As a fallback try converting the value to a tensor.
normalized_components.append(
ops.convert_to_tensor(t, name="component_%d" % i))
else:
if isinstance(spec, sparse_tensor.SparseTensorSpec):
normalized_components.append(
sparse_tensor.SparseTensor.from_value(t))
elif isinstance(spec, ragged_tensor.RaggedTensorSpec):
normalized_components.append(
ragged_tensor.convert_to_tensor_or_ragged_tensor(
t, name="component_%d" % i))
elif isinstance(spec, (tensor_array_ops.TensorArraySpec,
dataset_ops.DatasetSpec)):
normalized_components.append(t)
elif isinstance(spec, NoneTensorSpec):
normalized_components.append(NoneTensor())
elif isinstance(t, composite_tensor.CompositeTensor):
normalized_components.append(t)
else:
normalized_components.append(
ops.convert_to_tensor(t, name="component_%d" % i))
return nest.pack_sequence_as(element, normalized_components)
def get_classes(tensors):
"""Gets classes for a structure of tensors.
Args:
tensors: the tensor structure to get classes for.
Returns:
a structure matching the nested structure of `tensors`, containing
`tf.SparseTensor` at positions where `tensors` contains a sparse tensor and
`tf.Tensor` otherwise
"""
return nest.pack_sequence_as(tensors, [
sparse_tensor.SparseTensor
if isinstance(tensor, sparse_tensor.SparseTensor) else ops.Tensor
for tensor in nest.flatten(tensors)
])
def tpu_eval_step():
"""Generate the TPU graph."""
values = self.infeed_queue[0].generate_dequeue_op(tpu_device=0)
unflattened_inputs = data_nest.pack_sequence_as(self.feature_structure,
values)
features = unflattened_inputs["features"]
estimator_spec = model_fn(features, None, tf.estimator.ModeKeys.PREDICT,
params)
for k, v in six.iteritems(estimator_spec.predictions):
self.outfeed_names.append(k)
self.outfeed_tensors.append(v)
with tf.device(utils.device_for_tpu_core(self._get_host(0))):
outfeed_enqueue_ops = tpu.outfeed_enqueue_tuple(self.outfeed_tensors)
with tf.control_dependencies([outfeed_enqueue_ops]):
return tf.no_op()
def get_next(self, name=None):
"""Returns a nested structure of `tf.Tensor`s containing the next element.
Args:
name: (Optional.) A name for the created operation.
Returns:
A nested structure of `tf.Tensor` objects.
"""
return nest.pack_sequence_as(
self._output_types,
gen_dataset_ops.iterator_get_next(
self._iterator_resource,
output_types=nest.flatten(self._output_types),
output_shapes=nest.flatten(self._output_shapes),
name=name))
def serialize_sparse_tensors(tensors):
"""Serializes sparse tensors.
Args:
tensors: a tensor structure to serialize.
Returns:
`tensors` with any sparse tensors replaced by the their serialized version.
"""
ret = nest.pack_sequence_as(tensors, [
sparse_ops.serialize_sparse(tensor)
if isinstance(tensor, sparse_tensor.SparseTensor) else tensor
for tensor in nest.flatten(tensors)
])
return ret
def get_sparse_types(tensors):
"""Gets sparse types for a structure of tensors.
Args:
tensors: the tensor structure to get sparse types for.
Returns:
a structure matching the nested structure of `tensors`, containing
`SparseType` at positions where `tensors` contains a sparse tensor and
`None` otherwise
"""
return nest.pack_sequence_as(tensors, [
SparseType(tensor.dtype)
if isinstance(tensor, sparse_tensor.SparseTensor) else None
for tensor in nest.flatten(tensors)
])
def unwrap_sparse_types(types):
"""Unwraps sparse tensor types as `dtypes.string`.
Args:
types: a structure of types to unwrap.
Returns:
a structure matching the nested structure of `types`, containing
`dtypes.string` at positions where `types` contains a sparse tensor and
matching contents of `types` otherwise
"""
ret = nest.pack_sequence_as(types, [
dtypes.string if isinstance(ty, SparseType) else ty
for ty in nest.flatten(types)
])
return ret
