pyg-team/pytorch_geometric

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

Open

#4,543 opened on 2022幎4月26日

GitHub で芋る
 (23 comments) (6 reactions) (1 assignee)Python (3,514 forks)batch import
0 - Priority P0datafeaturehelp wanted

Repository metrics

Stars
 (19,985 stars)
PR merge metrics
 (平均マヌゞ 35d 1h) (30d で 14 merged PRs)

説明

🚀 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

コントリビュヌタヌガむド