Example for flax.jax_utils.prefetch_to_device with TFDS dataset

[davidshen84]

Hi,

I want to use flax.jax_utils.prefetch_to_device to preload data to my GPU. But I could not figure out how to work with a TFDS dataset.

The dataset is simple.

(train_ds,) = tfds.load("mnist", split=["train"])
train_ds = train_ds.repeat(...).shuffle(...).map(...).batch(batch_size).take(train_steps)
...
for step, batch in enumerate(train_ds.as_numpy_iterator()):
  ...

The above code works without using prefetch_to_device. But I cannot simply call prefetch_to_device(train_ds.as_numpy_iterator()) because prefetch_to_device requires the first dimension of the iterator to be the number of devices. I only have one GPU, so it expects the as_numpy_iterator to return an iterator in the shape of (1, batch_size, ...).

I cannot find a way to make as_numpy_iterator to return one more dimension. Maybe I need a new way to constract the training dataset? Anyway, it'd be very appreciated if I could get some code snippet or examples.

Thank you.

[n2cholas]

Have you tried calling .batch() twice, so that you have two batch dimensions. e.g. https://github.com/n2cholas/progan-flax/blob/main/src/main.py#L41-L43

[davidshen84]

It took me a while to understand how the batch dimensions were changed. I found it did not help with the training performance, at least not for a single GPU case.