[RFC] [DataLoader] Make GraphDataLoader support feature prefetching and offline sampling
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Make GraphDataLoader support feature prefetching and offline sampling.
## Motivation
@nv-dlasalle mentioned an offline sampling scenario, where the users might perform minibatch sampling beforehand and store the sampled subgraph structures on an external source. If the features are stored together with the sampled subgraphs, then the functionality can be well supported by `GraphDataLoader`. However, often times the users will store node/edge data on a separate source, rather than storing together with the sampled subgraphs. Otherwise, the same node data will duplicate across multiple sampled subgraphs, wasting a lot of space.
Currently DGL's `DataLoader` (not `GraphDataLoader`) does not directly slice the node/edge data during sampling; instead, it places some `LazyFeatures` marked for prefetching in the worker processes, and the actual slicing is done on the master process. We would like to migrate that functionality to `GraphDataLoader` as well.
## Alternatives
#4431 proposed a unification of `DataLoader` and `GraphDataLoader` by changing the `DataLoader` interface to accept `GraphSource` and `FeatureSource` objects.
## Pitch
My proposed UX looks like this:
```python
### User code 1 for node classification
# Say if this is already an iterable of (input nodes, output nodes, blocks)
# triplets stored as binary files on the HDFS. Only the indices are stored
# in the blocks, not the features.
# The users can implement it like any kind of PyTorch iterable-style or
# map-style dataset.
class HDFSGraphDataset(object):
def __init__(self, paths):
self.paths = paths
def __getitem__(self, key):
path = self.paths[key]
input_nodes, output_nodes, blocks = get_data_from_hdfs(path)
return input_nodes, output_nodes, blocks
def __len__(self):
return len(self.paths)
dataset = HDFSGraphDataset(...)
# SetBlockLazyFeatures takes in an iterable, and transforms the iterable by
# assuming that the last element is a list of blocks, and puts the input,
# output and edge LazyFeatures there.
dataset = dgl.dataloading.transforms.SetBlockListLazyFeatures(
dataset,
prefetch_node_feats=['feat'],
prefetch_labels=['label'],
prefetch_edge_feats=['weight'])
dataloader = GraphDataLoader(dataset, batch_size=None)
dataloader.attach_ndata('feat', FeatureStorage(...))
dataloader.attach_ndata('label', FeatureStorage(...))
dataloader.attach_edata('weight', FeatureStorage(...))
for input_nodes, output_nodes, blocks in dataloader:
pass
### User code 2 for link prediction
# Say if this is already an iterable of (input nodes, pair graph, negative
# pair graph, blocks) tuples stored as binary files on the HDFS. Only the
# indices are stored in the blocks, not the features.
class HDFSGraphDataset(object):
def __init__(self, paths):
self.paths = paths
def __getitem__(self, key):
path = self.paths[key]
input_nodes, pair_graph, neg_pair_graph, blocks = get_data_from_hdfs(path)
return input_nodes, pair_graph, neg_pair_graph, blocks
def __len__(self):
return len(self.paths)
dataset = HDFSGraphDataset(...)
dataset = dgl.dataloading.transforms.SetBlockListLazyFeatures(
dataset,
prefetch_node_feats=['feat'],
prefetch_edge_feats=['weight'])
# SetPairGraphLazyFeatures assumes that the second element of the iterable
# is a pair graph from edge classification/link prediction, and puts the
# edge labels there.
dataset = dgl.dataloading.transforms.SetPairGraphLazyFeatures(
dataset,
prefetch_labels=['label'])
dataloader = GraphDataLoader(dataset, batch_size=None)
dataloader.attach_ndata('feat', FeatureStorage(...))
dataloader.attach_edata('label', FeatureStorage(...))
dataloader.attach_edata('weight', FeatureStorage(...))
for input_nodes, pair_graph, neg_pair_graph, blocks in dataloader:
pass
```
There are three parts to implement:
For one, methods `attach_ndata`, `attach_edata` and `attach_data` as described in #4443 should be added to `GraphDataLoader` as well. Since there is no concept of "original graph", I think decoupling `FeatureStorage` from `GraphStorage` is a necessity.
For the other, DGL need to provide some mechanisms that insert "LazyFeatures" into the sampled subgraphs yielded by the graph dataset, so that the master process can retrieve them later. These mechanisms can be implemented via *transform*ing a map-style or an iterable-style graph dataset into another dataset:
* `dgl.dataloading.transforms.SetBlockListLazyFeatures` will assume that the last element of the yielded result is a list of blocks. It assigns LazyFeatures so that the master process can retrieve them from the associated FeatureStorage later.
* There should also be a `dgl.dataloading.transforms.SetSubgraphLazyFeatures` for subgraph sampling of course.
* `dgl.dataloading.transforms.SetPairGraphLazyFeatures` is another transform, that assigns the LazyFeatures to the second element instead, assuming it a pair graph.
The transforms can be implemented as an iterable/map wrapper over another iterable/map. The users can add in other transforms in a similar fashion to do whatever they want. For instance, if the sampled subgraphs are stored as plain indices instead of in DGL's format, then they can implement a transform that creates the graphs/blocks from those indices.
## Additional context
The idea of transforms are in fact identical to the concept of PyTorch DataPipes. I'm also thinking of refactoring the entire DataLoader implementation with DataPipes but obviously that takes a lot more effort (may even have performance regression), so this proposal could also be a smaller step towards the goal.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.