Hi, looking forward to the new optimizations!
However, I use cifar10 to easily showing what difficulty I have.
In fact, I am working on a PR for image captioning in Flax, and it will use the feature extractor of the ViT model. This feature extraction will convert RGB image into float32 3D array. And when I want to index into a preprocessed dataset (that uses this feature extractor) to get a batch, it takes quite long time.
The dataset I have is a collection of local images. It would be great if I can make it work without this slowness of processing.
For example, to get a batch of 256 elements in a preprocessed datasets, it takes about 5 seconds.
If you have a clear idea, that would be great. Otherwise, I understand you have other priorities