Get batch indices from NeighborLoader

[itsdaniele]

Hello!

I am trying to use NeighborLoader with certain batch size. The code is:

train_loader = NeighborLoader(
          data,
          num_neighbors=[5, 10, 10],
          num_workers=32,
          input_nodes=data["train_mask"],
          batch_size=8
)

However, whenever I retrieve a batch, I also need to densify it in a single tensor by calling to_dense_batch().

The issue is that the batches that are coming out of the dataloader don't have the batch indices tensor that allows me to assign each node to its own batch.

That is, if I do data = next(iter(train_loader)), then the data object doesn't have the data.batch tensor that I need to pass to to_dense_batch().

Am I missing something here?

[rusty1s]

The NeighborLoader will construct a single computation graph for all nodes within your mini-batch, it will not create a isolated subgraphs for each central node. I suggest to take a look into the ShaDowKHopSampler in PyG which does exactly this :)