Hi !
There are two optimizations that you can use:

Use the Array3D feature type. Otherwise it will consider your arrays to be lists of arbitrary sizes, and it takes some time to collate them again into one numpy array.

features = Features({
    **train_ds.features,
    "image": Array3D(dtype="int32", shape=(128, 128, 3))
})
train_ds = train_ds.map(
    preprocess_function,
    batched=True,
    batch_size=16,
    num_proc=2,
    features=features
)


Use the ‘numpy’ format. This way when accessing examples they will already be numpy arrays, and you will also benefit from the end-to-end zero-copy array read from Arrow to numpy.

train_ds = train_ds.with_format("numpy")

It becomes 100 times faster when querying the data loader !
Here are my results without the optimizations:
Average times_1 in data_loader: 7.690094709396362

And then, with these optimizations:
Average times_1 in data_loader: 0.0582084059715271