tf.contrib.data.bucket_by_sequence_length fails for nested Dataset element

[lhlmgr]

Hello everyone,

I just tried the new function to group variable length inputs for the dataset API, namely: tf.contrib.data.bucket_by_sequence_length, for a small Estimator-Model.

I implemented the input_fn such that it returns a dataset, where each element is a tuple (feature-dict, label). However, when I run it, I get following exception:

Traceback (most recent call last):
...
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 960, in apply
dataset = transformation_func(self)
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 198, in _apply_fn
window_size_func=window_size_fn))
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 960, in apply
dataset = transformation_func(self)
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 90, in _apply_fn
window_size_func)
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 239, in init
self._make_key_func(key_func, input_dataset)
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 289, in _make_key_func
self._key_func.add_to_graph(ops.get_default_graph())
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/function.py", line 488, in add_to_graph
self._create_definition_if_needed()
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/function.py", line 321, in _create_definition_if_needed
self._create_definition_if_needed_impl()
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/function.py", line 338, in _create_definition_if_needed_impl
outputs = self._func(*inputs)
File "/home/leo/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 279, in tf_key_func
ret = key_func(*nested_args)
TypeError: element_to_bucket_id() takes 1 positional argument but 2 were given

Here is a link to the function.

Here is a code snipped to reproduce the error:

import tensorflow as tf

def input_fn():
  def generator():
    text = [[1, 2, 3],
            [3, 4, 5, 6, 7],
            [1, 2],
            [8, 9, 0, 2, 3]]
    label = [1, 2, 1, 2]

    for x, y in zip(text, label):
      yield (x, y)

  dataset = tf.data.Dataset.from_generator(generator=generator,
                                           output_shapes=(tf.TensorShape([None]), tf.TensorShape([])),
                                           output_types=(tf.int32, tf.int32))

  dataset = dataset.map(parse_example)
  dataset = dataset.apply(tf.contrib.data.bucket_by_sequence_length(element_length_func=element_length_fn,
                                                                    bucket_batch_sizes=[2, 2, 2],
                                                                    bucket_boundaries=[0, 8],
                                                                    pad_to_bucket_boundary=False))

  return dataset

def parse_example(x, y):
  return dict(
    x=x
  ), y

def element_length_fn(element):
  features, label = element
  return tf.shape(features["x"])[0]

if __name__ == '__main__':
  with tf.Session() as sess:
    dataset = input_fn()
    iter = dataset.make_one_shot_iterator()

    print(sess.run(iter.get_next()))

My Env-Specs are logged in: tf_env.txt

Thanks in advance!

[mrry]

@rsepassi Can you take a look, please?

/cc @ebrevdo FYI.

[rsepassi]

Yes, looking into it.

…

On Thu, Mar 22, 2018 at 9:02 PM Derek Murray ***@***.***> wrote: @rsepassi <https://github.com/rsepassi> Can you take a look, please? /cc @ebrevdo <https://github.com/ebrevdo> FYI. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#17932 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/ABEGW-jtSsBu9OcEFUz8J3yEIi4AdEo5ks5thEmIgaJpZM4S3plo> .

[rsepassi]

This is consistent with the behavior of all callables in tf.data with tuple inputs. The function _should_unpack_args is used to determine whether inputs to a user-supplied callable should be unpacked (i.e. called as fn(*args)) and is True if the inputs are a tuple. If the elements of a Dataset are tuples, as they are in the example in this thread, then callables must expect the arguments to be passed unpacked.

For example, here's a Dataset where the elements are tuples and a call to map fails when the map_fn expects only 1 element.

    def data_gen():
      text = [[1, 2, 3], [3, 4, 5, 6, 7], [1, 2], [8, 9, 0, 2, 3]]
      label = [1, 2, 1, 2]
      for x, y in zip(text, label):
        yield (x, y)

    dataset = tf.data.Dataset.from_generator(
        generator=data_gen,
        output_shapes=(tensor_shape.TensorShape([None]), tensor_shape.TensorShape([])),
        output_types=(dtypes.int32, dtypes.int32))

    # Fails
    dataset.map(lambda el: el)

We may want to update this behavior but that's where it stands right now.
Fix right now would be to use a dict going in and apply a map to turn into a tuple afterwards.

@mrry, should we look into updating the behavior of Dataset callables?

[mrry]

@rsepassi Right, that example should fail, and we can't change this even if we wanted to. It works if you instead do dataset.map(lambda x, y: (x, y)).

I think there might be a missing _should_unpack_args() in the bucket_by_sequence_length code, because making the equivalent change to @lhlmgr's program doesn't fix things, i.e.:

def element_length_fn(features, label):
  return tf.shape(features["x"])[0]

...still yields the same error:

TypeError: element_to_bucket_id() takes exactly 1 argument (2 given)

[rsepassi]

Yup, found the issue. Have a fix out. bucket_by_sequence_length was not correctly handling tupleized elements.