How to sample batches for multiple graphs with NeighborLoader?

[ys-zong]

Hi, thanks for the nice work!

I'm trying to use NeighborLoader to sample batches for multiple graphs. For one graph, I have

from torch_geometric.loader import NeighborLoader
from torch_geometric.data import Data
 
mydata = Data(x = ..., edge_index = ...)
train_loader = NeighborLoader(mydata, num_neighbors=[10, 10, 10], batch_size=50)

Now I have multiple graphs and I want to sample batches like the above code (each subgraph contains neighbors within one graph, a batch contains multiple subgraphs from different graphs). For the normal dataloader, we can use

data_list = [Data(...), ..., Data(...)]
loader = DataLoader(data_list, batch_size=32)

I'm wondering if it also works for NeighborLoader and how do I check if it correctly samples in the way I want? In addition, after sampling, how can I know which graph is the current sampled data from?

Thank you for your help in advance!

[rusty1s]

In order to apply NeighborLoader to multiple graphs, one how to group all graphs into a supergraph first via Batch.from_data_list(data_list). The batch vector after sampling then denotes the assignment of nodes to their respective graph. Hope this helps.

[ys-zong]

Great, it works! Thank you! Just to confirm, after loading the Batch to NeighborLoader, I have my_batch = next(iter(loader)). The values in my_batch.batch indicate the assignments of the respective graph?

[rusty1s]

Yes, that is correct.