Dataset slow during model training

## Describe the bug
While migrating towards 🤗 Datasets, I encountered an odd performance degradation: training suddenly slows down dramatically. I train with an image dataset using Keras and execute a `to_tf_dataset` just before training.

First, I have optimized my dataset following https://discuss.huggingface.co/t/solved-image-dataset-seems-slow-for-larger-image-size/10960/6, which actually improved the situation from what I had before but did not completely solve it.

Second, I saved and loaded my dataset using `tf.data.experimental.save` and `tf.data.experimental.load` before training (for which I would have expected no performance change). However, I ended up with the performance I had before tinkering with 🤗 Datasets.

Any idea what's the reason for this and how to speed-up training with 🤗 Datasets?

## Steps to reproduce the bug
```python
# Sample code to reproduce the bug

from datasets import load_dataset
import os

dataset_dir = "./dataset"
prep_dataset_dir = "./prepdataset"
model_dir = "./model"

# Load Data
dataset = load_dataset("Lehrig/Monkey-Species-Collection", "downsized")
def read_image_file(example):
    with open(example["image"].filename, "rb") as f:
        example["image"] = {"bytes": f.read()}
        return example
dataset = dataset.map(read_image_file)
dataset.save_to_disk(dataset_dir)

# Preprocess
from datasets import (
    Array3D,
    DatasetDict,
    Features,
    load_from_disk,
    Sequence,
    Value
)
import numpy as np
from transformers import ImageFeatureExtractionMixin

dataset = load_from_disk(dataset_dir)

num_classes = dataset["train"].features["label"].num_classes
one_hot_matrix = np.eye(num_classes)
feature_extractor = ImageFeatureExtractionMixin()

def to_pixels(image):
    image = feature_extractor.resize(image, size=size)
    image = feature_extractor.to_numpy_array(image, channel_first=False)
    image = image / 255.0
    return image

def process(examples):
    examples["pixel_values"] = [
        to_pixels(image) for image in examples["image"]
    ]
    examples["label"] = [
        one_hot_matrix[label] for label in examples["label"]
    ]
    return examples

features = Features({
    "pixel_values": Array3D(dtype="float32", shape=(size, size, 3)),
    "label": Sequence(feature=Value(dtype="int32"), length=num_classes)
})

prep_dataset = dataset.map(
    process,
    remove_columns=["image"],
    batched=True,
    batch_size=batch_size,
    num_proc=2,
    features=features,
)

prep_dataset = prep_dataset.with_format("numpy")

# Split
train_dev_dataset = prep_dataset['test'].train_test_split(
    test_size=test_size,
    shuffle=True,
    seed=seed
)

train_dev_test_dataset = DatasetDict({
    'train': train_dev_dataset['train'],
    'dev': train_dev_dataset['test'],
    'test': prep_dataset['test'],
})

train_dev_test_dataset.save_to_disk(prep_dataset_dir)

# Train Model
import datetime
import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras.applications import InceptionV3
from tensorflow.keras.layers import Dense, Dropout, GlobalAveragePooling2D, BatchNormalization
from tensorflow.keras.callbacks import ReduceLROnPlateau, ModelCheckpoint, EarlyStopping
from transformers import DefaultDataCollator

dataset = load_from_disk(prep_data_dir)

data_collator = DefaultDataCollator(return_tensors="tf")

train_dataset = dataset["train"].to_tf_dataset(
    columns=['pixel_values'],
    label_cols=['label'],
    shuffle=True,
    batch_size=batch_size,
    collate_fn=data_collator
)

validation_dataset = dataset["dev"].to_tf_dataset(
    columns=['pixel_values'],
    label_cols=['label'],
    shuffle=False,
    batch_size=batch_size,
    collate_fn=data_collator
)

print(f'{datetime.datetime.now()} - Saving Data')
tf.data.experimental.save(train_dataset, model_dir+"/train")
tf.data.experimental.save(validation_dataset, model_dir+"/val")

print(f'{datetime.datetime.now()} - Loading Data')
train_dataset = tf.data.experimental.load(model_dir+"/train")
validation_dataset = tf.data.experimental.load(model_dir+"/val")

shape = np.shape(dataset["train"][0]["pixel_values"])
backbone = InceptionV3(
    include_top=False,
    weights='imagenet',
    input_shape=shape
)

for layer in backbone.layers:
    layer.trainable = False

model = Sequential()
model.add(backbone)
model.add(GlobalAveragePooling2D())
model.add(Dense(128, activation='relu'))
model.add(BatchNormalization())
model.add(Dropout(0.3))
model.add(Dense(64, activation='relu'))
model.add(BatchNormalization())
model.add(Dropout(0.3))
model.add(Dense(10, activation='softmax'))

model.compile(
    optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

print(model.summary())

earlyStopping = EarlyStopping(
    monitor='val_loss',
    patience=10,
    verbose=0,
    mode='min'
)

mcp_save = ModelCheckpoint(
    f'{model_dir}/best_model.hdf5',
    save_best_only=True,
    monitor='val_loss',
    mode='min'
)

reduce_lr_loss = ReduceLROnPlateau(
    monitor='val_loss',
    factor=0.1,
    patience=7,
    verbose=1,
    min_delta=0.0001,
    mode='min'
)

hist = model.fit(
    train_dataset,
    epochs=epochs,
    validation_data=validation_dataset,
    callbacks=[earlyStopping, mcp_save, reduce_lr_loss]
)
```

## Expected results
Same performance when training without my "save/load hack" or a good explanation/recommendation about the issue.

## Actual results
Performance slower without my "save/load hack".

**Epoch Breakdown (without my "save/load hack"):**
- Epoch 1/10
41s 2s/step - loss: 1.6302 - accuracy: 0.5048 - val_loss: 1.4713 - val_accuracy: 0.3273 - lr: 0.0010
- Epoch 2/10
32s 2s/step - loss: 0.5357 - accuracy: 0.8510 - val_loss: 1.0447 - val_accuracy: 0.5818 - lr: 0.0010
- Epoch 3/10
36s 3s/step - loss: 0.3547 - accuracy: 0.9231 - val_loss: 0.6245 - val_accuracy: 0.7091 - lr: 0.0010
- Epoch 4/10
36s 3s/step - loss: 0.2721 - accuracy: 0.9231 - val_loss: 0.3395 - val_accuracy: 0.9091 - lr: 0.0010
- Epoch 5/10
32s 2s/step - loss: 0.1676 - accuracy: 0.9856 - val_loss: 0.2187 - val_accuracy: 0.9636 - lr: 0.0010
- Epoch 6/10
42s 3s/step - loss: 0.2066 - accuracy: 0.9615 - val_loss: 0.1635 - val_accuracy: 0.9636 - lr: 0.0010
- Epoch 7/10
32s 2s/step - loss: 0.1814 - accuracy: 0.9423 - val_loss: 0.1418 - val_accuracy: 0.9636 - lr: 0.0010
- Epoch 8/10
32s 2s/step - loss: 0.1301 - accuracy: 0.9856 - val_loss: 0.1388 - val_accuracy: 0.9818 - lr: 0.0010
- Epoch 9/10
loss: 0.1102 - accuracy: 0.9856 - val_loss: 0.1185 - val_accuracy: 0.9818 - lr: 0.0010
- Epoch 10/10
32s 2s/step - loss: 0.1013 - accuracy: 0.9808 - val_loss: 0.0978 - val_accuracy: 0.9818 - lr: 0.0010



**Epoch Breakdown (with my "save/load hack"):**
- Epoch 1/10
13s 625ms/step - loss: 3.0478 - accuracy: 0.1146 - val_loss: 2.3061 - val_accuracy: 0.0727 - lr: 0.0010
- Epoch 2/10
0s 80ms/step - loss: 2.3105 - accuracy: 0.2656 - val_loss: 2.3085 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 3/10
0s 77ms/step - loss: 1.8608 - accuracy: 0.3542 - val_loss: 2.3130 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 4/10
1s 98ms/step - loss: 1.8677 - accuracy: 0.3750 - val_loss: 2.3157 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 5/10
1s 204ms/step - loss: 1.5561 - accuracy: 0.4583 - val_loss: 2.3049 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 6/10
1s 210ms/step - loss: 1.4657 - accuracy: 0.4896 - val_loss: 2.2944 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 7/10
1s 205ms/step - loss: 1.4018 - accuracy: 0.5312 - val_loss: 2.2917 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 8/10
1s 207ms/step - loss: 1.2370 - accuracy: 0.5729 - val_loss: 2.2814 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 9/10
1s 214ms/step - loss: 1.1190 - accuracy: 0.6250 - val_loss: 2.2733 - val_accuracy: 0.0909 - lr: 0.0010
- Epoch 10/10
1s 207ms/step - loss: 1.1484 - accuracy: 0.6302 - val_loss: 2.2624 - val_accuracy: 0.0909 - lr: 0.0010

## Environment info
- `datasets` version: 2.2.2
- Platform: Linux-4.18.0-305.45.1.el8_4.ppc64le-ppc64le-with-glibc2.17
- Python version: 3.8.13
- PyArrow version: 7.0.0
- Pandas version: 1.4.2
- TensorFlow: 2.8.0
- GPU (used during training): Tesla V100-SXM2-32GB


Comments:
Hi ! cc @Rocketknight1 maybe you know better ?

I'm not too familiar with `tf.data.experimental.save`. Note that `datasets` uses memory mapping, so depending on your hardware and the disk you are using you can expect performance differences with a dataset loaded in RAM

Hi @lehrig, I suspect what's happening here is that our `to_tf_dataset()` method has some performance issues when streaming samples. This is usually not a problem, but they become apparent when streaming a vision dataset into a very small vision model, which will need a lot of sample throughput to saturate the GPU.

When you save a `tf.data.Dataset` with `tf.data.experimental.save`, all of the samples from the dataset (which are, in this case, batches of images), are saved to disk. When you load this saved dataset, you're effectively bypassing `to_tf_dataset()` entirely, which alleviates this performance bottleneck.

`to_tf_dataset()` is something we're actively working on overhauling right now - particularly for image datasets, we want to make it possible to access the underlying images with `tf.data` without going through the current layer of indirection with `Arrow`, which should massively improve simplicity and performance. 

However, if you just want this to work quickly but without needing your save/load hack, my advice would be to simply load the dataset into memory if it's small enough to fit. Since all your samples have the same dimensions, you can do this simply with:

```
dataset = load_from_disk(prep_data_dir)
dataset = dataset.with_format("numpy")
data_in_memory = dataset[:]
```

Then you can simply do something like:

```
model.fit(data_in_memory["pixel_values"], data_in_memory["labels"])
```

Thanks for the information! 

I have now updated the training code like so:

```
dataset = load_from_disk(prep_data_dir)
train_dataset = dataset["train"][:]
validation_dataset = dataset["dev"][:]

...

model.fit(
    train_dataset["pixel_values"],
    train_dataset["label"],
    epochs=epochs,
    validation_data=(
        validation_dataset["pixel_values"],
        validation_dataset["label"]
    ),
    callbacks=[earlyStopping, mcp_save, reduce_lr_loss]
)
```

- Creating the in-memory dataset is quite quick
- But: There is now a long wait (~4-5 Minutes) before the training starts (why?)
- And: Training times have improved but the very first epoch leaves me wondering why it takes so long (why?)

**Epoch Breakdown:**
- Epoch 1/10
78s 12s/step - loss: 3.1307 - accuracy: 0.0737 - val_loss: 2.2827 - val_accuracy: 0.1273 - lr: 0.0010
- Epoch 2/10
1s 168ms/step - loss: 2.3616 - accuracy: 0.2350 - val_loss: 2.2679 - val_accuracy: 0.2182 - lr: 0.0010
- Epoch 3/10
1s 189ms/step - loss: 2.0221 - accuracy: 0.3180 - val_loss: 2.2670 - val_accuracy: 0.1818 - lr: 0.0010
- Epoch 4/10
0s 67ms/step - loss: 1.8895 - accuracy: 0.3548 - val_loss: 2.2771 - val_accuracy: 0.1273 - lr: 0.0010
- Epoch 5/10
0s 67ms/step - loss: 1.7846 - accuracy: 0.3963 - val_loss: 2.2860 - val_accuracy: 0.1455 - lr: 0.0010
- Epoch 6/10
0s 65ms/step - loss: 1.5946 - accuracy: 0.4516 - val_loss: 2.2938 - val_accuracy: 0.1636 - lr: 0.0010
- Epoch 7/10
0s 63ms/step - loss: 1.4217 - accuracy: 0.5115 - val_loss: 2.2968 - val_accuracy: 0.2182 - lr: 0.0010
- Epoch 8/10
0s 67ms/step - loss: 1.3089 - accuracy: 0.5438 - val_loss: 2.2842 - val_accuracy: 0.2182 - lr: 0.0010
- Epoch 9/10
1s 184ms/step - loss: 1.2480 - accuracy: 0.5806 - val_loss: 2.2652 - val_accuracy: 0.1818 - lr: 0.0010
- Epoch 10/10
0s 65ms/step - loss: 1.2699 - accuracy: 0.5622 - val_loss: 2.2670 - val_accuracy: 0.2000 - lr: 0.0010



Regarding the new long ~5 min. wait introduced by the in-memory dataset update: this might be causing it? https://datascience.stackexchange.com/questions/33364/why-model-fit-generator-in-keras-is-taking-so-much-time-even-before-picking-the

For now, my save/load hack is still more performant, even though having more boiler-plate code :/ 

That 5 minute wait is quite surprising! I don't have a good explanation for why it's happening, but it can't be an issue with `datasets` or `tf.data` because you're just fitting directly on Numpy arrays at this point. All I can suggest is seeing if you can isolate the issue - for example, does fitting on a smaller dataset containing only 10% of the original data reduce the wait? This might indicate the delay is caused by your data being copied or converted somehow. Alternatively, you could try removing things like callbacks and seeing if you could isolate the issue there.