Use prefetch() after batching // image_dataset.py

[Frightera]

I think it is better use prefetch() after batching. In a pipeline we would want the next batch ready (t+1) while processing current batch (t).

[qlzh727]

Adding @jsimsa for the performance best practice here.

From https://www.tensorflow.org/guide/data_performance#prefetching, I didn't see any suggestion about whether the prefetch should be applied before batching or not.

The current approach will prefetch unbatched data (with autotune), unless the prefetched data can't fill the next batch, then I don't see any big difference here.

[Frightera]

@qlzh727

Quoting from official documentation:

Most dataset input pipelines should end with a call to prefetch. This allows later elements to be prepared while the current element is being processed. This often improves latency and throughput, at the cost of using additional memory to store prefetched elements.

I think it is better to prefetch the batches rather than single elements, so we can have next batch(es) ready while processing current batch in the training process.

[qlzh727]

Thanks for the reference, let's wait for some inputs from tf.data side.

[qlzh727]

Chatted with @wilsingosti offline for this issue: (copied from the chat)

Prefetch is useful in 2 cases:

1. If the upstream of Prefetch has variance, inserting Prefetch can help to buffer enough elements to reduce variance

2. if there is a long sequence of synchronous transformations, eg. sequential map, shuffle, batch, inserting a Prefetch somewhere in the sequence can break the sequence such that the original transformations can be pipelined.

In this particular case, since the input of batch is a ParallelMap, inserting a Prefetch between batch and ParallelMap will not help case 2 because ParallelMap is an asynchronous op (which means that it has its own buffer). Inserting it after batch can help especially if the output of batch is another synchronous transform.

[qlzh727]

Based on the assessment above, I am approving this PR. Thanks @Frightera for the contribution.

[Frightera]

Thanks @qlzh727, those statements were also helpful.