This is a question about grain, the python library for data ingestion. https://google-grain.readthedocs.io/en/latest/index.html In the context of reinforcement learning I have a data pre-processing step where I have match replays as individual files, pre-process into JAX arrays, and then process. I would be nice to have a library that understand the logic of transforming lists of JAX arrays and bundle them together into sequences for training recurrent networks, and into batches for gradient descent. I think I have almost all components. I've implemented a ReplayDataSource which subclasses grain.RandomAccessDataSource which opens and pre-processes the match file into memory. I've implemented a Sampler which returns contiguous indices to build sequences. With one data source I can have all I need with (schematically) source = ReplayDataSource( ... ) transformations = [grain.Batch(batch_size)] sampler = MySampler( ... ) data_loader = grain.DataLoader( data_source = source, sampler = sampler, operations = transformations, shard_options=grain.NoSharding(), ) This will add 2 outer dimensions to my dataset, one for sequences and another for batches. Great! There's one element missing. If I have say 1000 replay files and I can't load all to memory, how do I tell the DataLoader to instantiate load ReplayDataSources into memory and get rid of them as needed, as I iterate on the elements of the DataLoader? I've tried source = [source1, source2] on the DataLoader args but that didn't work. I've try making a MapDataset out of the data sources, but I don't know how to compose those either. Same with the IterDataset. The only grain functionality I really care about is sequences and batches. How do I create a dataset compose out of make replay files and only open those as needed? Continue reading...
