Partition Large Graph for Testing

[melsaa]

I have a large graph (12M nodes, 66M edges) and my task is to do link prediction with GNN. I successfully trained the model with GraphSAINT, however, the learned embedding is too big to be stored in memory.

Is there any partition/clustering algorithm that preserve all the edges so that I can use it during testing?

Thanks

[rusty1s]

Sadly no, and this is one of the disadvantages of subgraph-sampling since you typically want to do inference on the full graph. There is a clever workaround though, which allows you to create node-embeddings on CPU in a layer-wise fashion, i.e., it will create all node embeddings for the first layer, and then uses those to create node embeddings for the second layer, and so on. Here is an example on how to do so.

[melsaa]

Thanks for the answer. So, from what I understand GPU memory still has to be large enough to store all node features and one layer embedding. Is it possible to do link prediction for a new node where all the links to the original graph are unknown?

[rusty1s]

That is not really true, since we save the input features and node embeddings on CPU rather than GPU. Only the current mini-batch will be put on GPU.

Regarding your second question: This is indeed a bit tricky, since you can only infer connections for the new nodes based on its initial node features. However, in theory, this should be doable via a classic link prediction approach.