Something about NeighborLoader's mini-batch sampling and in-batch splitting.

[skdbsxir]

Hello :)

First, thanks for contributing those excellent works and many examples.

Currently I started studying GNNs and their methodologies(e.g. data handling in GNNs), I'm just curious about NeighborLoader.

Using "Planetoid/cora" dataset, I'm just checking how mini-batching via NeighborLoader works.
(Not multiple graphs, just single graph case.)

1. Something about num_neighbors parameter.

• Following some lectures and tutorials, I understood NeighborLoader makes each mini-batch by making computation graph using seed nodes and their randomly sampled neighbors.
• And number of seed nodes inside mini-batches are same as batch_size, and number of randomly sampled neighbors are same as num_neighbors, if I understood correctly.
• I just wanted to watch how many neighbors are sampled to each seed node(aka wanted to watch node's degree), I just set num_neighbors = [-1]. But it seems not all neighbors are selected.
  • I added n_id to data and checked each node's degree, and simply made mini-batch by NeighborLoader with num_neighbors = [-1], batch_size = 2, and compared them as below code.

import networkx as nx
import pandas as pd

from torch_geometric.datasets import Planetoid
from torch_geometric.loader import NeighborLoader
from torch_geometric.seed import seed_everything
from torch_geometric.utils import to_networkx

seed_everything(62)

dataset = Planetoid(root='dataset', name='cora')

data = dataset[0]

graph = to_networkx(data)
degree_list = list(graph.degree)
df = pd.DataFrame(degree_list, columns=['n_id', 'degree'])
# print(df.head())
"""
   n_id  degree
0     0       6
1     1       6
2     2      10
3     3       2
4     4      10
"""
loader = NeighborLoader(
    data = data,
    num_neighbors = [-1],
    batch_size = 2,
)

## check batch
subgraph_list = []
for batch in loader:
    subgraph_list.append(batch)

# print('Total # of batches : ', len(subgraph_list)) # 1354 (2708 / batch_size=2)

## check 1st batch -> 2 computation graphs with 2 seed nodes (n_id = 0, 1)
# print(subgraph_list[0].n_id.tolist())
# print(subgraph_list[0].edge_index.numpy())
"""
Total # of batches :  1354
[0, 1, 1862, 633, 2582, 2, 654, 652]
[[2 3 4 5 6 7]
 [0 0 0 1 1 1]]
"""

• It seems two results have some difference between original node's degree and computation graph with all neighbors.
  • Original degree of n_id=0 and n_id=1 is 6 and 6 each, but (mini-batch's) computation graph selected only 3 nodes I confirmed.
  • Why is it happens? Just I misunderstood something?

2. Selecting individual computation graph in single mini-batch.

• For some reason, I just want to split each computation graph (constructed by selected seed node and randomly selected neighbors) inside single mini-batch.
  • For example, I just want to split like below.

[[2 3 4 5 6 7]
 [0 0 0 1 1 1]]

to

[[2 3 4]
 [0 0 0]]

[[5 6 7]
 [1 1 1]]

• Is there any way to split each of them like above in PyG? Or just split them by myself?

Thank you in advance for any! :)

[rusty1s]

I think you need to specify to_networkx(..., to_undirected=True) in order to get the correct node degree from networkx.

[skdbsxir]

Gosh, I thought cora citation data was directed graph. Checked it with data.is_undirected() was True. Thanks so much for kind reply!

Additionally, is there any kind of methods in PyG to do such thing as my question 2?

I think setting NeighborLoader's parameter disjoint=True and using torch_geometric.utils.unbatch_edge_index to each mini-batch will returns something like I wanted to do, but result seems little bit different what I expected.

from torch_geometric.datasets import Planetoid
from torch_geometric.loader import NeighborLoader
from torch_geometric.utils import unbatch_edge_index

dataset = Planetoid(root='dataset', name='cora')

data = dataset[0]
data.n_id = torch.arange(data.num_nodes)

loader = NeighborLoader(
    data = data,
    num_neighbors = [-1],
    batch_size = 4,
    disjoint=True # each seed node will create its own disjoint subgraph
)

## check batch
subgraph_list = []
for batch in loader:
    subgraph_list.append(batch)

# print('Total # of batches : ', len(subgraph_list))
# print(subgraph_list[0])
# print('Selected nodes in this mini-batch', subgraph_list[0].n_id.tolist())
# print('Where each nodes are in batch : ' , subgraph_list[0].batch.tolist())
# print(subgraph_list[0].edge_index.numpy())
"""
Total # of batches :  677
Data(x=[16, 1433], edge_index=[2, 12], y=[16], train_mask=[16], val_mask=[16], test_mask=[16], n_id=[16], batch=[16], input_id=[4], batch_size=4)
Selected nodes in this mini-batch [0, 1, 2, 3, 1862, 633, 2582, 2, 654, 652, 1666, 1, 1454, 1986, 332, 2544]
Where each nodes are in batch :  [0, 1, 2, 3, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2, 3]
[[ 4  5  6  7  8  9 10 11 12 13 14 15]
 [ 0  0  0  1  1  1  2  2  2  2  2  3]]
"""

## Batch splitting? 
batch_list_in_minibatch = unbatch_edge_index(subgraph_list[0].edge_index, batch=subgraph_list[0].batch)
# print(batch_list_in_minibatch)

"""
(tensor([[4, 5, 6],
        [0, 0, 0]]), tensor([[ 3,  4,  5],
        [-3, -3, -3]]), tensor([[ 2,  3,  4,  5,  6],
        [-6, -6, -6, -6, -6]]), tensor([[  1],
        [-11]]))
"""

I thought subgraph_list[0].edge_index would split like my question 2, but why is there minus value in result? Did I miscatched something?

[rusty1s]

Yes, you would need to set disjoint=True. I think unbatch_edge_index does not necessarily do what one expects since the batch vector is not sorted. Alternatively, you can do:

edge_batch = batch[batch.edge_index[0]]
for i in range(batch_size)
    edge_index_for_i = data.edge_index[:, edge_batch == i]

[skdbsxir]

Thanks a lot for your kind reply!

I checked out your reply's code, wrote a code like below, and found out what I wanted.

sample_batch = subgraph_list[0]

# print(sample_batch.edge_index.numpy())
"""
[[ 4  5  6  7  8  9 10 11 12 13 14 15]
 [ 0  0  0  1  1  1  2  2  2  2  2  3]]
"""

sample_graph_edge_index = sample_batch.batch[sample_batch.edge_index[0]]

# print(sample_graph_edge_index.tolist())
"""
[0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2, 3]
"""
## BATCH_SIZE is 4
for i in range(BATCH_SIZE):
    edge_index_for_i = sample_batch.edge_index[:, sample_graph_edge_index == i]
    # print(edge_index_for_i)
"""
tensor([[4, 5, 6],
        [0, 0, 0]])
tensor([[7, 8, 9],
        [1, 1, 1]])
tensor([[10, 11, 12, 13, 14],
        [ 2,  2,  2,  2,  2]])
tensor([[15],
        [ 3]])
"""

And I could found out above's code can do the right thing what I wanted, when using not only as num_neighbors=[-1](1-layer), also as num_neighbors=[2, 3](2-layer).

num_neighbors=[2,3]'s result was like below.

# sample_batch.edge_index.numpy()
[[ 4  5  6  7  8  9 10  5 11 12  4  0 13 14 15  1  1  2  2 16 17  3]
 [ 0  0  1  1  2  2  3  4  4  4  5  5  5  6  6  6  7  8  9  9  9 10]]

# sample_graph_edge_index.tolist()
[0, 0, 1, 1, 2, 2, 3, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3]

# edge_index_for_i.numpy()
[[ 4  5  5 11 12  4  0 13]
 [ 0  0  4  4  4  5  5  5]]
[[ 6  7 14 15  1  1]
 [ 1  1  6  6  6  7]]
[[ 8  9  2  2 16 17]
 [ 2  2  8  9  9  9]]
[[10  3]
 [ 3 10]]

So, each elements in sample_graph_edge_index(in your code is batch[batch.edge_index[0]]) means it corresponds to that id's computation graph, and by using edge_index's slicing I can get right results (correctly partitioned computation graphs).
(Correct me if I misunderstood above's flow)