Skip to content

feat(models): Support loading graphs from disk #1150

Description

@Sooph-knmi

Is your feature request related to a problem? Please describe.

Currently, the graph provider supports either static graphs or dynamically created graphs.
When scaling up to a large amount of (large-sized) graphs with either multi-dataset or multi-domain, using static graphs can cause out of memory exceptions, and dynamically creating graphs causes and extreme slow-down in training.

Describe the solution you'd like

Pre-generating graphs and loading them from disk would solve both problems by reducing memory footprint with minimal time increase. Initial tests show 250 300MB graphs can be loaded and offloaded in 182.0 ms (4 workers, pinned memory).

The proposed solution would be to add a graph provider FileGraphProvider allowing users to specify a folder with pre-generated graphs that would be loaded dynamically.

path = self.paths[idx]
graph = torch.load(path, weights_only=False, map_location="cpu")

Describe alternatives you've considered

An alternative solution would be to load all graphs in CPU memory. However, there remains a limit to how many graphs can be stored (experiments show a maximum of about 50 300MB graphs). Sharing memory across GPUs would increase the limit to 200 graphs, but no more.

Additional context

Use cases:

  • Training with large (high-resolution) graphs
  • Scaling up the amount of graphs (with multi-domain or multi-dataset training)

Organisation

KNMI

Metadata

Metadata

Labels

Projects

Status
Now In Progress

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions