JAX efficient data loading and processing

[Rasmuskh]

Hi,
I am wondering about how I should deal with data loading and processing in JAX in a local single GPU setup?
So far I have been doing the preprocessing (e.g. normalization) on CPU, and then applied random cropping and random flips on minibatches on the GPU in my train_step function.

However, my approach has some issues:

• Doing the preprocessing on CPU is extremely slow for large datasets (like ImageNet).
• Not having the data augmentation build into the dataloader makes it necessary to check which augmentations to apply inside my train_step function, which feels clumbsy/inelegant.

So far I have not been able to run the Tensorflow and JAX part of the code on the GPU at the same time. Could this have something to do with how VRAM is reserved by tensoflow and JAX?

So my main question is how to do both preprocessing and data augmentation on GPU in an elegant way?

Note: My question is closely related to this unanswered question from 2022: #13339

[alonfnt]

Not really an answer but there's the unreleased grain library: https://github.com/google/grain
which will give efficient data loading to GPU.

[Rasmuskh]

Thanks, this looks interesting!
I found a mention of this library in the tensorflow datasets docs: https://www.tensorflow.org/datasets/tfless_tfds#coming_soon_use_with_jax

But there doesn't seem to be any documentation or examples of how to use grain yet.

[braniii]

The TF article includes now grain:

https://www.tensorflow.org/datasets/data_source#use_with_jax