GraphSAGE not working with huge data with NeighborLoader function

[Chaitanyasrini]

I am currently working on a project involving the GraphSAGE algorithm, implemented using PyTorch Geometric, and I've encountered a couple of issues that I am hoping you could help me with.

Issue: Currently I am using NeighborLoader function to implement minibatching in GraphSAGE, the processing time for GraphSAGE becomes considerably longer. However, the extent of the slowdown seems disproportionate, and the jupyter notebook gets stuck. I am currently stuck in Neighbor loader as it is taking too much time to process data. Can anyone help me here:
Below is the dataset size and the parameters used in NeighborLoader and GraphSAGE methods given in PyG library
Parameters:
No. of Nodes:312321916
No. of edges: 661874770

Neighbor Loader Parameters:

input_nodes=None,
num_neighbors=[5,2]
batch_size=128,
replace=False,
shuffle=True

GraphSAGE Parameters:
in_channels=8,
hidden_channels=128,
num_layers=2

[wsad1]

The number of iterations in each epoch is going be around (2.4 million = 312321916/ 128). Depending on your hardware you could

1. increase the batch size. Depends on wether you have enough GPU to store larger graph.
2. Don't train on all nodes. Use input_nodes to train on just a subset of nodes. Many times training on a subset of nodes works as well as training on all nodes.
3. Try playing with num_workers parameter to see if you get speed ups.