pyg-team/pytorch_geometric

Interface for datasets that are too large to use `InMemoryDataset`

Open

#4543 aperta il 26 apr 2022

Vedi su GitHub
 (23 commenti) (6 reazioni) (1 assegnatario)Python (3514 fork)batch import
0 - Priority P0datafeaturehelp wanted

Metriche repository

Star
 (19.985 star)
Metriche merge PR
 (Merge medio 35g 1h) (14 PR mergiate in 30 g)

Descrizione

🚀 The feature, motivation and pitch

There are several examples of datasets for molecular property prediction where each individual graph example easily fits in memory but there are too many examples to fit within the InMemoryDataset interface. One solution is to save each example in its own .pt file but this introduces a significant filesystem overhead to access each example.

A better solution is to partition the data such that there are many graphs serialised within a single .pt file. The number of graphs can be considered a chunk_size parameter which is independent from the training batch_size. This ChunkedDataset interface would be expected to scale to as large a dataset as desired while avoiding the significant overhead of having one graph per file.

The design idea is roughly:

  • ChunkedDataset inherits from the PyG Dataset interface
  • Accepts a chunk_size argument
  • Has an abstract method process_chunk that accepts a list of data objects that can be processed and saved as a single .pt file.

Other considerations:

  • The training batch size should not depend on the chunk_size so the dataset
  • ChunkedDataset should support splitting to read from parallel workers as well as random shuffling

Alternatives

No response

Additional context

No response

Guida contributor