Large Dataset Best Practice

[LTMeyer]

Let me first thank you for all the great work behind pytorch_geometric.

I'm dealing with several temporal graphs that represent numerical simulations of physical systems. I have between tens and thousands of graph topology (that are actually meshes). For each topology, I have hundreds of time steps. Eventually, this yields to a dataset of several Gb, that cannot fit fully in memory.

The documentation indicates how to handle such large datasets with the torch_geometric.data.Dataset class. I decided to process and save a .pt file for each topology. Given an index, the get method will figure which topology and time step it corresponds to. It will then read the relevant .pt file and access the appropriate time step. Finally, it will return a Data object with all necessary inputs and targets.

The training on such dataset is naturally slower compared to InMemoryDataset. Especially, when the batch increases the loading of the data is much slower. However, I was wondering what were the best practices to make the access of the data as fast as possible. What are the general rules of thumb in this case? What would be the best splits in terms of saved files, topology, and time step? A processed file per time step would lead to smaller but numerous files. Grouped several topology and time step would lead to fewer but bigger files. Is there also any consideration to have regarding the number of workers, or is it simply the more the faster? How much processed information must be saved in .pt or computed on the fly?

[rusty1s]

I think in general you want to avoid saving too many small files, as I/O will become a bottleneck. However, you already may find success in increasing the number of workers num_workers to load files in parallel. I think the ideal number of workers needs to be determined empirically and depend on the given CPU.

The splitting also depends on your model. If your model makes use of temporal information, I think it's best to store all timesteps into a single Data object. In addition, one can also think about storing small mini-batches into individual files and create larger ones within the DataLoader. However, I do not have an example script to show-case this.

[LTMeyer]

I think in general you want to avoid saving too many small files, as I/O will become a bottleneck. However, you already may find success in increasing the number of workers num_workers to load files in parallel. I think the ideal number of workers needs to be determined empirically and depend on the given CPU.

I guess the answer of the following question will depend on the hardware too, but is there any rule of thumb regarding the size the files should be to avoid I/O bottleneck?

The splitting also depends on your model. If your model makes use of temporal information, I think it's best to store all timesteps into a single Data object.

Yes it makes totally sense. Unfortunately, this is not the case. One time step corresponds to one input of my model.

In addition, one can also think about storing small mini-batches into individual files and create larger ones within the DataLoader. However, I do not have an example script to show-case this.

If I understand clearly, here you're suggesting to already store mini-batches in the processed file? The dataset would then be shuffled before being saved. For example, file 0 would now include the samples [graph_2, t_6], [graph_10, t_74] ... to give arbitrary numbers where the first one corresponds to the topology and the second one the time steps. Rather than file 0 being [graph_0, t_0], [graph_0, t_1] ... [graph_0, t_n_time_steps]. This way, to generate a batch the Dataloader will have to read less files (reducing the I/O) bottleneck.

[rusty1s]

I sadly cannot give you a rule of thumb. This is dependent on your file size, batch size and I/O speed. I think this is a general problem in Deep Learning, and this blog post might be of interest to you.

Yes, your understanding about storing small mini-batches is correct and I think this might be a good method to approach this problem. Let me know if you can achieve any benefits in loading that way :)

[LTMeyer]

Thank you for these pointers. I think it is indeed the direction to go.
The webdataset library seems to be helpful. The documentation is not mature yet, but the functionalities are there. I will give it a try. In the meantime we may consider this discussion as resolved.

Thank you very much for your help.

[rusty1s]

Thanks for the pointer to webdataset which I did not know yet. Let me know if you are interested in contributing an example on how to combine webdataset and PyTorch Geometric :)

[LTMeyer]

No worries, it was actually mentioned in the blog post you pointed to.

I would be happy to provide such an example that combines webdataset and pytorch_geometric. However, I should first get some time to familiarise with the library in order to provide a fully working and relevant example. Besides, there is a discussion on integrating webdataset directly into pytorch's code base. But this may have little impacts on the example.

In any case, if I find something relevant worth to be added I will make a PR with link to this conversation.