random and augmentation in datasets

[pcasl]

I am trying to do an on-fly-augmentation using the dataset.

train_dataset = (grain.MapDataset.source(source)
                     .map(augment())
                     .shuffle(seed=412)
                     .repeat(num_epochs=config.n_epochs)
                     .batch(batch_size=config.batch_size, drop_remainder=True)
                    )
train_dataset = grain.experimental.ThreadPrefetchIterDataset(train_dataset, prefetch_buffer_size=64)
train_dataset = train_dataset.map(jax.device_put)

The augment() requires random seeds. But the train_dataset will be prefeched in parallel threads, and I am not sure how to pass the random seed. Any suggestions?

ps. I am also not sure that the shuffle was done correctly. Please correct me if I'm wrong.

Thanks a lot!

[iindyk]

I think train_dataset.random_map(augment, seed =42) is what you're looking for

https://google-grain.readthedocs.io/en/latest/grain.dataset.html#grain.MapDataset.random_map

the augment will get an element and a numpy RNG that you should use to do the random augmentation. It will be seeded deterministically and grain will ensure that the dataset is augmented differently depending on the element index and epoch