Repository: tensorflow
Issue: tf.data batching slows down on Windows
Labels: stat:awaiting tensorflower,type:bug,comp:data,TF 2.8

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Source

binary

### Tensorflow Version

tf 2.8

### Custom Code

Yes

### OS Platform and Distribution

Windows 10

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

CUDA 11.2 cuDNN 8.1.0

### GPU model and memory

Nvidia A40 48GB, 128GB System RAM

### Current Behaviour?


While training a keras model using a [custom tensorflow dataset](https://www.tensorflow.org/datasets/add_dataset), the training steps slow down once the remaining sample count in the first epoch is smaller than the shuffle buffer size. The slowdown carries over to the following epochs and takes longer and longer.

I would expect the opposite behavior since all remaining samples should already be in the buffer and do not need to be loaded anymore. I have tested it on multiple Windows systems and this slowdown occurred on all of them. However, it does not happen on Linux-based systems.

I then used the tensorboard profiling plugin to investigate what is causing the slowdown. As you can see here, it seems to be an input-related problem:
![profiler_1](https://user-images.githubusercontent.com/1509163/166112234-b795f11e-934f-4631-86b1-c8d79e1e1448.png)

From the input operations you can see that `Iterator::Root::Prefetch::BatchV2` takes the most time (this is also the case when removing the prefetching):
![profiler_2](https://user-images.githubusercontent.com/1509163/166112243-798a33a9-b329-4103-9d24-b0b70306411d.png)

The slowdown can also be seen in the trace viewer:
![profiler_3](https://user-images.githubusercontent.com/1509163/166112328-9e637031-1166-4ac6-afef-47a8a534bf9b.png)

Here is a more detailed comparison of `BatchV2` durations from different steps:
![profiler_4](https://user-images.githubusercontent.com/1509163/166112335-119bb08e-7cd2-4a4b-a149-0de1580332a4.png)
![profiler_5](https://user-images.githubusercontent.com/1509163/166112342-dd0ea8a4-9d79-4ec3-adc9-87ac133e118e.png)


Here is the [profiling_data.zip](https://github.com/tensorflow/tensorflow/files/8597248/profiling_data.zip).



### Standalone code to reproduce the issue

In order to reproduce and profile this issue the following code can be used. When running the `main.py` for the first time it will generate the dummy dataset which takes around 20 minutes to write 24GB of dummy data.

**dummy_dataset.py**
```python
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds


class DummyDataset(tfds.core.GeneratorBasedBuilder):
    VERSION = tfds.core.Version('1.0.0')
    RELEASE_NOTES = {
        '1.0.0': 'Initial release.',
    }

    def _info(self) -> tfds.core.DatasetInfo:
        return tfds.core.DatasetInfo(
            builder=self,
            features=tfds.features.FeaturesDict({
                'audio': tfds.features.Audio(shape=(16000,), dtype=tf.float32),
                'label': tfds.features.Tensor(shape=(2,), dtype=tf.float32),
            }),
            supervised_keys=('audio', 'label')
        )

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        return {
            'train': self._generate_examples(),
        }

    def _generate_examples(self):
        for i in range(400000): # If sample count is increased, issue starts later (tested with 1M+)
            yield i, {
                'audio': np.random.sample(16000).astype(dtype=np.float32),
                'label': tf.one_hot(np.random.choice(2), depth=2).numpy()
            }
```



**main.py**
```python
from dummy_dataset import DummyDataset
import tensorflow_datasets as tfds
import tensorflow as tf
import os
import datetime
from tensorflow import keras

# Configuration
shuffle_buffer = 300000 # Uses ~30GB RAM, can be lowered to 150000 if only 16GB RAM available
batch_size = 512
epochs = 3

optimizer = keras.optimizers.Adam(learning_rate=0.001)
losses = ['categorical_crossentropy']
metrics = [keras.metrics.CategoricalAccuracy()]

folder_name = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

# Prepare dataset (takes ~20min for the first time to generate ~24GB of dummy data)
ds_train, ds_info = tfds.load(
    name='DummyDataset',
    split='train',
    as_supervised=True,
    with_info=True,
    shuffle_files=True,
)
ds_train = ds_train.shuffle(shuffle_buffer)
ds_train = ds_train.batch(batch_size)
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

# Configure profiling
sample_count = ds_info.splits['train'].num_examples
slowdown_start = int((sample_count - shuffle_buffer) / batch_size)
profile_stop = slowdown_start + 150

callbacks = [keras.callbacks.TensorBoard(
    log_dir=os.path.join(os.getcwd(), 'tensorboard', folder_name),
    profile_batch=[slowdown_start, profile_stop]
)]

# Build and train model
input_layer = keras.layers.Input(shape=(16000,), dtype=tf.float32, name='audio_input')
output_layer = keras.layers.Dense(2, activation='softmax', name='prediction')(input_layer)
model = keras.Model(inputs=input_layer, outputs=output_layer, name='dummy_model')

model.compile(optimizer=optimizer, loss=losses, metrics=metrics)
model.summary()

print(f'Expected slowdown to start at batch {slowdown_start}, profiling batches {slowdown_start}-{profile_stop}')

model.fit(x=ds_train, epochs=epochs, verbose=1, callbacks=callbacks)
```

**requirements.txt**
```
tensorflow==2.8.0
tensorflow-datasets==4.5.2
tensorboard-plugin-profile
numpy
```


### Relevant log output

_No response_</details>