Sampling In NodeNeighborhood Loader

[akul-goyal]

I am using a NodeNeighborhood loader to sample nodes from my graph. Currently, I have the batch size sent to 32, the number of neighbors = [-1], and shuffle to False. I assume the loader iterates in batches of 32 through the node set. My understanding of the node sampling is that only the representations of those 32 nodes selected should be updated. However, when iterating through the loader I am seeing more than 32 nodes as part of the initial feature vector. Are all the nodes within the feature vector being updated or just the original 32 nodes that were sampled? Moreover, how can you select the nodes that were to be sampled from the resulting node embedding tensor outputted from the model given the shape of the tensor holds more than the number of nodes you sampled?

[rusty1s]

Thanks for your interest and sorry for the late reply. In your case, the resulting mini-batch does hold the original nodes as well as all of its direct neighbors. You can select the original nodes from your feature matrix/embedding matrix by simply slicing the first 32 nodes from it viax[:32] (or x[:data.batch_size] if you want to be agnostic to batch sizes).

Let me know if that answers your question :)

[akul-goyal]

Thank you for the reply! Do the original node's direct neighbors also get updated during the message passing phase of the GNN?

[rusty1s]

It depends on what you mean with getting updated. They will be used to update their intermediate representations and the representations of their neighbors as well. Importantly, you should only read out the node embeddings of the input nodes though.