[Roadmap] Advanced Graph Sampling Routines 🚀

[motwaniravi12]

🚀 The feature, motivation and pitch

Different GNN applications require customized sampling to improve the accuracy and/or performance. We propose the following advanced sampling routines with detailed benchmarking reports.

• Multi-threading support: All sampling operations are single-threaded. We can expect additional speed-ups when providing multi-threading implementation:
  • Multi-threaded neighbor_sample implementation
  • Multi-threaded hetero_neighbor_sample implementation (Hetero neighbor sampler multithreading pyg-lib#215)
  • Introduce benchmarks in pyg-lib and/or PyG
• Bi-directional sampling: Allow the option to return a bidirectional graph as part of sampling (e.g., for allowing GNNs with more layers than sampled hops)
  • Add SubgraphType with bidirectional sampling support (1/2) #7199
  • Add SubgraphType with bidirectional sampling support (2/2) #7200
  • Introduce bidirectional sampling benchmarks
• Hierarchical Neighborhood sampling: Hierarchical Neighborhood Sampling is extending classical Neighborhood Sampling by collecting additional information about number of sampled nodes and edges per each hop. That information can be used in special features like Hierarchical Graph Adjacency Matrix implemented by trim_to_layer functionality which helps to boost the performance.
  • Adding option to return info about sampled graph pyg-lib#197
  • Adding trim_to_layer utility function #6661
  • SparseTensor support for trim_to_layer #7089
  • Adds an example for Hierarchical Sampling #7244
  • Add automatic trimming example via a separate HGAM layer and add an example
  • Write documentation/tutorial
• Weighted Sampling: Allow passing a weight vector to neighbor_sample to sample based on
  • Allow passing a weight vector to neighbor_sample to sample based on importance scores. Reference DGL implementation (see prob argument).
  • Integrate into PyG
• Subgraph sampling
  • Initial version in pyg-lib, see here
  • Benchmarking efforts
  • Integrate into pyg_lib.neighbor_sample in case directed=False

[xiachenrui]

Hi, directed=False parameter in NeighborLoader() will raise 'RuntimeError: Undirected subgraphs not yet supported' yet. Will this be fixed anytime soon ？

Best,
Chen-Rui

[rusty1s]

Yes, this is bullet point 3 under "Subgraph sampling".

[xiachenrui]

Happy to hear this ! I can't wait to try it.

[rusty1s]

This routine is implemented in torch-sparse, so you can either do torch_geometric.typing.WITH_PYG_LIB=False or use the subgraph_type="bidirectional" argument in the meantime.

[xiachenrui]

OK, I notice the subgraph_type parameter is avaliable in the latest version of docments. But the latest release (2.3.1) do not include this update. I resolve this by install from github.

Thanks a lot ! @rusty1s

[mfbalin]

Would there be any interest in a contribution of the following sampler from DGL?https://docs.dgl.ai/generated/dgl.dataloading.LaborSampler.html#dgl.dataloading.LaborSampler
It is basically a combination of Neighbor and Layer Sampling methods. DGL currently has a working CPU, GPU, and GPU over UVA implementation. Can be used as a drop-in replacement to neighbor_sample as labor_sample.

[rusty1s]

Thanks for the pointer. What would be needed to integrate this into PyG?

[mfbalin]

I used the pcg32 random number generator as the PRNG, given random seed z and vertex id t, r_t = PRNG(z, t) in the paper. So for the CPU implementation, there will be a need for either implementing or including https://github.com/imneme/pcg-cpp as a submodule. Other fast PRNGs could also be used if they are fast at generating random numbers as a function of the random seed z and a given vertex id t.

An initial CPU contribution should be a fairly small PR.

[mfbalin]

My PR is here: pyg-team/pyg-lib#242.