def test_type_to_tf_dtypes_and_shapes_with_two_level_tuple(self):
dtypes, shapes = type_utils.type_to_tf_dtypes_and_shapes(
collections.OrderedDict([
('a', tf.bool),
('b',
collections.OrderedDict([
('c', tf.float32),
('d', (tf.int32, [20])),
])),
]))
test.assert_nested_struct_eq(dtypes, {
'a': tf.bool,
'b': {
'c': tf.float32,
'd': tf.int32
}
})
test.assert_nested_struct_eq(
shapes, {
'a': tf.TensorShape([]),
'b': {
'c': tf.TensorShape([]),
'd': tf.TensorShape([20])
}
})
def _assert_binding_matches_type_and_value(self, binding, type_spec, val,
graph):
"""Asserts that 'bindings' matches the given type, value, and graph."""
self.assertIsInstance(binding, pb.TensorFlow.Binding)
self.assertIsInstance(type_spec, computation_types.Type)
binding_oneof = binding.WhichOneof('binding')
if binding_oneof == 'tensor':
self.assertTrue(tf.contrib.framework.is_tensor(val))
if not isinstance(val, tf.Variable):
# We insert a read_value() op for Variables, which produces
# a name we don't control. Otherwise, names should match:
self.assertEqual(binding.tensor.tensor_name, val.name)
self.assertIsInstance(type_spec, computation_types.TensorType)
self.assertEqual(type_spec.dtype, val.dtype.base_dtype)
self.assertEqual(repr(type_spec.shape), repr(val.shape))
elif binding_oneof == 'sequence':
self.assertIsInstance(val,
graph_utils.DATASET_REPRESENTATION_TYPES)
sequence_oneof = binding.sequence.WhichOneof('binding')
if sequence_oneof == 'iterator_string_handle_name':
# TODO(b/129956296): Eventually delete this deprecated code path.
handle = graph.get_tensor_by_name(
binding.sequence.iterator_string_handle_name)
self.assertIn(str(handle.op.type),
['Placeholder', 'IteratorToStringHandle'])
self.assertEqual(handle.dtype, tf.string)
else:
self.assertEqual(sequence_oneof, 'variant_tensor_name')
variant_tensor = graph.get_tensor_by_name(
binding.sequence.variant_tensor_name)
op = str(variant_tensor.op.type)
self.assertTrue((op == 'Placeholder') or ('Dataset' in op))
self.assertEqual(variant_tensor.dtype, tf.variant)
self.assertIsInstance(type_spec, computation_types.SequenceType)
output_dtypes, output_shapes = (
type_utils.type_to_tf_dtypes_and_shapes(type_spec.element))
test.assert_nested_struct_eq(
tf.compat.v1.data.get_output_types(val), output_dtypes)
test.assert_nested_struct_eq(
tf.compat.v1.data.get_output_shapes(val), output_shapes)
elif binding_oneof == 'tuple':
self.assertIsInstance(type_spec, computation_types.NamedTupleType)
if not isinstance(val,
(list, tuple, anonymous_tuple.AnonymousTuple)):
self.assertIsInstance(val, dict)
if isinstance(val, collections.OrderedDict):
val = list(val.values())
else:
val = [v for _, v in sorted(six.iteritems(val))]
for idx, e in enumerate(anonymous_tuple.to_elements(type_spec)):
self._assert_binding_matches_type_and_value(
binding.tuple.element[idx], e[1], val[idx], graph)
else:
self.fail('Unknown binding.')
def test_type_to_tf_structure_without_names(self):
type_spec = computation_types.to_type((tf.bool, tf.int32))
dtypes, shapes = type_utils.type_to_tf_dtypes_and_shapes(type_spec)
structure = type_utils.type_to_tf_structure(type_spec)
with tf.Graph().as_default():
ds = tf.data.experimental.from_variant(
tf.placeholder(tf.variant, shape=[]), structure=structure)
ds_dtypes = tf.compat.v1.data.get_output_types(ds)
ds_shapes = tf.compat.v1.data.get_output_shapes(ds)
test.assert_nested_struct_eq(ds_dtypes, dtypes)
test.assert_nested_struct_eq(ds_shapes, shapes)
def test_type_to_tf_dtypes_and_shapes_with_tensor_triple(self):
dtypes, shapes = type_utils.type_to_tf_dtypes_and_shapes(
collections.OrderedDict([('a', (tf.int32, [5])), ('b', tf.bool),
('c', (tf.float32, [3]))]))
test.assert_nested_struct_eq(dtypes, {
'a': tf.int32,
'b': tf.bool,
'c': tf.float32
})
test.assert_nested_struct_eq(shapes, {
'a': tf.TensorShape([5]),
'b': tf.TensorShape([]),
'c': tf.TensorShape([3])
})
def make_dataset_from_string_handle(handle, type_spec):
"""Constructs a `tf.data.Dataset` from a string handle tensor and type spec.
Args:
handle: The tensor that represents the string handle.
type_spec: The type spec of elements of the data set, either an instance of
`types.Type` or something convertible to it.
Returns:
A corresponding instance of `tf.data.Dataset`.
"""
# TODO(b/129956296): Eventually delete this deprecated code path.
type_spec = computation_types.to_type(type_spec)
tf_dtypes, shapes = type_utils.type_to_tf_dtypes_and_shapes(type_spec)
def make(handle=handle, tf_dtypes=tf_dtypes, shapes=shapes):
"""An embedded no-argument function that constructs the data set on-demand.
This is invoked by `OneShotDataset` on each access to the data set argument
passed to the body of the TF computation to ensure that the iterators and
tje map are constructed in the appropriate context (e.g., in a defun).
Args:
handle: Captured from the local (above).
tf_dtypes: Captured from the local (above).
shapes: Captured from the local (above).
Returns:
An instance of `tf.data.Dataset`.
"""
with handle.graph.as_default():
it = tf.data.Iterator.from_string_handle(handle, tf_dtypes, shapes)
# In order to convert an iterator into something that looks like a data
# set, we create a dummy data set that consists of an infinite sequence
# of zeroes, and filter it through a map function that invokes
# 'it.get_next()' for each of those zeroes.
# TODO(b/113112108): Possibly replace this with something more canonical
# if and when we can find adequate support for abstractly defined data
# sets (at the moment of this writing it does not appear to exist yet).
return tf.data.Dataset.range(1).repeat().map(
lambda _: it.get_next())
# NOTE: To revert to the old behavior, simply invoke `make()` here directly.
return OneShotDataset(make, type_spec)
def test_type_to_tf_structure_with_names(self):
type_spec = computation_types.to_type(
collections.OrderedDict([
('a', tf.bool),
('b',
collections.OrderedDict([
('c', tf.float32),
('d', (tf.int32, [20])),
])),
]))
dtypes, shapes = type_utils.type_to_tf_dtypes_and_shapes(type_spec)
structure = type_utils.type_to_tf_structure(type_spec)
with tf.Graph().as_default():
ds = tf.data.experimental.from_variant(
tf.placeholder(tf.variant, shape=[]), structure=structure)
ds_dtypes = tf.compat.v1.data.get_output_types(ds)
ds_shapes = tf.compat.v1.data.get_output_shapes(ds)
test.assert_nested_struct_eq(ds_dtypes, dtypes)
test.assert_nested_struct_eq(ds_shapes, shapes)
def test_type_to_tf_dtypes_and_shapes_with_int_vector(self):
dtypes, shapes = type_utils.type_to_tf_dtypes_and_shapes(
(tf.int32, [10]))
test.assert_nested_struct_eq(dtypes, tf.int32)
test.assert_nested_struct_eq(shapes, tf.TensorShape([10]))
def make_data_set_from_elements(graph, elements, element_type):
"""Creates a `tf.data.Dataset` in `graph` from explicitly listed `elements`.
NOTE: The underlying implementation attempts to use the
`tf.data.Dataset.from_tensor_slices() method to build the data set quickly,
but this doesn't always work. The typical scenario where it breaks is one
with data set being composed of unequal batches. Typically, only the last
batch is odd, so on the first attempt, we try to construct two data sets,
one from all elements but the last one, and one from the last element, then
concatenate the two. In the unlikely case that this fails (e.g., because
all data set elements are batches of unequal sizes), we revert to the slow,
but reliable method of constructing data sets from singleton elements, and
then concatenating them all.
Args:
graph: The graph in which to construct the `tf.data.Dataset`.
elements: A list of elements.
element_type: The type of elements.
Returns:
The constructed `tf.data.Dataset` instance.
Raises:
TypeError: If element types do not match `element_type`.
ValueError: If the elements are of incompatible types and shapes.
"""
py_typecheck.check_type(graph, tf.Graph)
py_typecheck.check_type(elements, list)
element_type = computation_types.to_type(element_type)
py_typecheck.check_type(element_type, computation_types.Type)
def _make(element_subset):
structure = make_empty_list_structure_for_element_type_spec(
element_type)
for el in element_subset:
append_to_list_structure_for_element_type_spec(
structure, el, element_type)
tensor_slices = to_tensor_slices_from_list_structure_for_element_type_spec(
structure, element_type)
return tf.data.Dataset.from_tensor_slices(tensor_slices)
output_types, output_shapes = (
type_utils.type_to_tf_dtypes_and_shapes(element_type))
with graph.as_default():
if not elements:
# Just return an empty data set.
# TODO(b/124517334): Remove from_generator when the best option for
# creating empty dataset with specific type signature is identified.
ds = tf.data.Dataset.from_generator((lambda: ()),
output_types=output_types,
output_shapes=output_shapes)
elif len(elements) == 1:
ds = _make(elements)
else:
try:
# It is common for the last element to be a batch of a size different
# from all the preceding batches. With this in mind, we proactively
# single out the last element (optimizing for the common case).
ds = _make(elements[0:-1]).concatenate(_make(elements[-1:]))
except ValueError:
# In case elements beyond just the last one are of unequal shapes, we
# may have failed (the most likely cause), so fall back onto the slow
# process of constructing and joining data sets from singletons. Not
# optimizing this for now, as it's very unlikely in scenarios
# we're targeting.
ds = None
for i in range(len(elements)):
singleton_ds = _make(elements[i:i + 1])
ds = singleton_ds if ds is None else ds.concatenate(
singleton_ds)
ds_element_type = type_utils.tf_dtypes_and_shapes_to_type(
ds.output_types, ds.output_shapes)
if not type_utils.is_assignable_from(element_type, ds_element_type):
raise TypeError(
'Failure during data set construction, expected elements of type {}, '
'but the constructed data set has elements of type {}.'.format(
str(element_type), str(ds_element_type)))
return ds
