Mini batch training with NeighborLoader

[skalyanb]

I am trying to do a large-scale GNN learning for node classification. I am using the NeighborLoader to create a mini-batch as follows.

loader = NeighborLoader(data,
                            num_neighbors=[-1],
                # Use a batch size of 128 for sampling training nodes of type paper
                            batch_size=32,
                            input_nodes=data.train_mask,
                            shuffle=True)
for batch in loader:
        print(batch.batch_size)
        print(batch.x.shape)

My understanding is that NeighborLoader will select 32 nodes with data.train_mask label true in each batch. For each of the nodes in a batch, it will generate the computational graph with 1-neighborhood. Is that what is happening in the NeighborLoader?

Now, I want to know the node indices that are selected in each batch. How do I extract that information from the batch?

Thanks in advance for your help.

[wsad1]

Your understanding of how NeighborLoader works sounds good to me.

To get original node id add a tensor to your data object that indicates the original node id. As explained in this answer.

Maybe we should add a short note on this to the documentation since multiple people have the same question.