How can I sample random neighbors from an edge list?

[HacksonJohn]

Hi, I am trying to sample some random neighbors for the target nodes. I follow the document and use NeighborLoader. However, It doesn't work as expected --- the target nodes I input are not useful at all.

Is there a way I can fix it? An example and code are provided below.

---

1. Suppose we have a simple edge_list and node feature (6 nodes).

from torch_geometric.data import Data
from torch_geometric import utils
from torch_geometric import loader

edge_index = torch.tensor([[0, 0, 0, 1, 1, 1, 2, 3, 3],
                           [1, 2, 3, 2, 3, 4, 3, 4, 5]], dtype=torch.long)
edge_index = utils.add_self_loops(edge_index)[0]
edge_index = utils.to_undirected(edge_index)

x = torch.eye(6) # node features

data = Data(edge_index=edge_index, x=x) # Create the PyG Data object

2. Then I want to sample 2 random neighbors of the target nodes. For example, suppose we want to generate 2 random neighbors for target nodes [1,3].

seed_nodes = torch.tensor([1,3],dtype=torch.long)  # target nodes

ll = loader.NeighborLoader(data,
                           num_neighbors=[3],)
                           input_nodes=seed_nodes)

3. However, the results are not what we want. When we print the results, it show as follows.

for l in ll: print(l.edge_index)

>>>
tensor([[0, 1, 2],
        [0, 0, 0]])
tensor([[1, 2, 3],
        [0, 0, 0]])

Since we want the random neighbors for 1 and 3, why it outputs edges like [0,0] and [0,2]? May I ask what I did wrong, and how can I fix it? Thank you!

[rusty1s]

It is important to note that neighborhood sampling will re-map the global node indices to local indices of the sampled subgraph. If you wanna restore this mapping, you can do

data.n_id = torch.arange(data.num_nodes)

loader = NeighborLoader(data, ...)

batch = next(iter(loader))
edge_index = batch.n_id[batch.edge_index]

[HacksonJohn]

Perfectly Solved. Thank you so much!