Discuss Batch Standards in TFIO with Keras

[BryanCutler]

Following the discussion on #366 batching can serve different purposes and optimizing for each is not always done the same way.

Previously, we thought we will reuse the same batch concept for serving two purposes:
Batch is a way to improve performance for record based files such as parquet, feature, etc, as we really don't want to "read one integer or one float32 at a time". This is a caching issue.

We also want to directly feed the dataset to tf.keras so batch concept in tf.keras is reused (number of samples populated in neural network). This is a model parameter.
But those two batch concepts are different. @BryanCutler I am wondering if it makes sense to split those things out. In other words, for 1) we should "read as much record as possible, as long as it fits the memory", and for 2) we should do a rebatch() to adjust the batch_size that has been feeded to tf.keras?

[yongtang]

To adds additional information, tf.data actually has a tf.data.Dataset.concatenate API which effectively solves the issue of multiple files:

a = Dataset.range(1, 4)  # ==> [ 1, 2, 3 ]
b = Dataset.range(4, 8)  # ==> [ 4, 5, 6, 7 ]
a.concatenate(b)  # ==> [ 1, 2, 3, 4, 5, 6, 7 ]

I think we only need to focus on single filename case (as multiple files could be processed at higher level with concatenate). After Dataset has been concatenated batch mode could reapply:

a.concatenate(b).batch(batch_size, drop_reminder)

# or

a.batch(batch_size, drop_reminder).concatenate(b.batch(batch_size, drop_reminder)

I do think rebatch could be a useful feature which is not available in tf.data.Dataset. rebatch could be emulated with:

d.unbatch().batch(batch_size) # <= d.rebatch(batch_size)

but I suspect d.unbatch().batch(batch_size) is not going to have a good performance.

[yongtang]

Added a PR #393 to introduce rebatch(batch_size) ~= unbatch().batch(batch_size).

@terrytangyuan @BryanCutler Here is my thinking on tf.keras batch vs cache batch issue:

1. We want to max out the cache batch (unless not fit into memory)
2. the tf.keras batch_size could always be placed as the last step before feeding too tf.keras (so that caching in the middle are not split into small chunks (of several float32 numbers).

So overall I think we could:

1. All Dataset API will not expose a batch_size, but internally Dataset always batched (up to the cache size).
2. We provide a rebatch(batch_size) so that user could call it before pass to tf.keras.

# Note:
#     - each individual record is [h, w, c]
#     - total records could fill into [n, h, w, c]
#     - out Dataset shape will be [c, h, w, c] where c <= n
class Dataset:
    def __init__(self):
        # use a large cache of c size
...
...

dateset = Dataset().map()....filter() # shape is [None, h, w, c]

mode.fit(dataset.rebatch(batch_size))

[BryanCutler]

This sounds like a good paradigm to me. Do you plan some way for the user to set the cache batch size or is it determined by the Dataset?

For the Arrow datasets, I'll have to think a bit more on the best way to handle this. Currently, records are read as batch, which is effectively the "cache batch" just not in the form of tensors. If this outputs the entire batch as tensors then it would be a second copy of all the batched data, which might not be good.

Still, it might be useful to cache the batch as tensors and release the original record batch. Then the user could call a rebatch(), like above, in the call to keras model.fit()

[yongtang]

@BryanCutler for internal cache size I am thinking about "try to fine tune automatically" unless user override (through **kwargs). In one PR for Parquet https://github.com/tensorflow/io/pull/384/files I am using a capacity to fine tune it. Currently hardcoded as 65536 but I plan to add more fine tuning later.

The capacity could be overridden in **kwargs.

[yongtang]

@BryanCutler Haven't spend much time to look into Arrow yet, though it looks like there are two different types: the feather file format and the arrow streaming format.

The file format probably fits the Tensor case we discussed here as in theory, we could easily distribute a list of cached tensors to different nodes to improve the performance by concurrently running ops on list of chunked (cached) tensors.

The streaming format, depending on if it is replayable or not, may not fit the file format handling.