RFC: Generalizing tf.data batching using windowing and reducers

[ewilderj]

This RFC will be open for comment until Friday, August 10th, 2018.

cc @jsimsa @mrry

Generalizing tf.data batching using windowing and reducers

Status	Proposed
Author(s)	Jiri Simsa (Google)
Sponsor	Derek Murray (Google)
Updated	2018-07-26

Objective

This proposal addresses the known limitations of the current tf.data batching API:

• it provides a mechanism for padded batching of sparse tensors
• it facilitates customization of batching logic (users can now express batching logic as a pure Python function)
• it enables application of different batching logic on different components

[bhack]

What about video analysis or video inpainting subfields? Have we some common shift a stride use cases for these?

[jsimsa]

The stride/shift uses cases for video analysis I am aware of are covered by tf.contrib.data.sliding_window_batch.

[carlthome]

What is expected to be passed into the init_fn?

[jsimsa]

It is the group key, which only makes sense if the reducer is used as part of the tf.contrib.data.group_by_reducer. We are still debating whether to have an unused init_fn argument in the non-keyed reducer or different interfaces for keyed and non-keyed reducers.

[carlthome]

Super!

Having control over batching directly in the Python API would be great, and it would let me solve tensorflow/tensorflow#20781 by making a tf.data.Reducer that sorts the tensor on the batch axis by some metadata like sequence id.

[Mistobaan]

It feels like consuming the dataset in batches should be pushed downstream in the API. something along the line of .get_next(batch_size=#) The reduce API is something you would apply to the content of the dataset and not to its packaging/consumption.

[jsimsa]

Thank you for your comment @Mistobaan. How would your account for the limitations that the current proposal aims to address? In particular, I am not sure how would .get_next(batch_size=#) enable users to express different types of batching (e.g. regular, padded, truncated, ...) or per-component batching? I can also imagine that giving the user program the flexibility to change the batch dimension of consecutive elements of the pipeline would hamper certain optimization (e.g. prefetching).

[carlthome]

Looking a little more closely at this I wonder if the return of tf.data.Dataset.reduce() shouldn't also be a tf.data.Dataset, because that seems like it's slowly becoming the expected object type in the high-level APIs such as tf.estimator and tf.keras, particularly with eager execution as:

for minibatch in dataset:
    model.train_on_batch(minibatch)

And with lazy:

estimator.train(lambda: dataset)

model.fit(dataset, steps=num_batches_in_dataset)

It will probably be confusing if only some of the Dataset methods are for chaining. Thoughts?

[jsimsa]

@carlthome I would actually argue that because of the integration with tf.estimator or tf.keras you would want the result to be a tensor as opposed to a dataset.

As illustrated by the examples in the proposal, one common use of tf.data.Dataset.reduce() is to reduce windows of tensors of a dataset to a single tensor. If the reduced element would be a dataset instead, we would end up with a dataset of datasets which is something that tf.estimator or tf.keras would not necessarily know how to handle. If you wanted to pass a single tensor produce by tf.data.Dataset.reduce() to tf.estimator or tf.keras, you can always use tf.data.Dataset.from_tensors to wrap it in a dataset.

[ageron]

Did you mean "To introduce the concept of tf.data.Reducer"?

[georgesterpu]

Hi @jsimsa
Is it possible to use the improved API to pad sequences differently than with constant values ?
The tf.pad transformation supports the "reflect" and "symmetric" padding modes. Is a PaddedBatchDataset currently limited to the "constant" padding mode ?

[jsimsa]

As Example 1 illustrates, the API allows you to write your custom padding logic, applying arbitrary map transformation (e.g. using tf.pad) to the elements to be padded.

In contrast, padded_batch is limited to either constant (or maximum) padding.

[georgesterpu]

This looks really great, @jsimsa !
Would you please help me with an example of a pad_fn that when applied on a WindowDataset produces symmetrically right padded sequences up to the length of longest one ? There are a few differences between the examples in the RFC and the current codebase, and can't get the snippet below to work:

import tensorflow as tf
import numpy as np

data = []
for _ in range(50):
    random_length = np.random.randint(low=10, high=100)
    data.append(tf.ones([random_length, 128]))

dataset = tf.data.Dataset.from_generator(lambda: data, output_types=tf.float32)


def pad_fn(sequences):
    return tf.pad(sequences, paddings=None, mode='SYMMETRIC')


dataset = dataset.window(4)
dataset = dataset.map(pad_fn)

for elem in dataset:
    print(elem)  # expected shape: [4, max_ts_batch, 128]