How does is_sorted work for neighborloader?

[davidlow]

I am confused how is_sorted works in pytorch geometric's neighborloader. According to the docs, if you the data's edge_index is sorted by column, you can use is_sorted=true to save time and memory. However, when I try it, I receive substantially different batch sizes.

I asked this question on stackexchange but I figured this would be a more appropriate place.

Here's an example:

import torch
import torch_geometric as pyg
from torch_geometric.datasets import KarateClub

data_kc = KarateClub()
print(data_kc.data)
print(data_kc.data.edge_index)

loader_f = pyg.loader.NeighborLoader(data_kc.data, num_neighbors=[2,2,2], batch_size=3, is_sorted=False, shuffle=False)
loader_t = pyg.loader.NeighborLoader(data_kc.data, num_neighbors=[2,2,2], batch_size=3, is_sorted=True, shuffle=False)

b_f = next(iter(loader_f))
b_t = next(iter(loader_t))

print("batch with is_sorted=False =", b_f)
print("batch with is_sorted=True  =",b_t)

Output:

Data(x=[34, 34], edge_index=[2, 156], y=[34], train_mask=[34])
tensor([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  1,
          1,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  3,
          3,  3,  3,  3,  3,  4,  4,  4,  5,  5,  5,  5,  6,  6,  6,  6,  7,  7,
          7,  7,  8,  8,  8,  8,  8,  9,  9, 10, 10, 10, 11, 12, 12, 13, 13, 13,
         13, 13, 14, 14, 15, 15, 16, 16, 17, 17, 18, 18, 19, 19, 19, 20, 20, 21,
         21, 22, 22, 23, 23, 23, 23, 23, 24, 24, 24, 25, 25, 25, 26, 26, 27, 27,
         27, 27, 28, 28, 28, 29, 29, 29, 29, 30, 30, 30, 30, 31, 31, 31, 31, 31,
         31, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33,
         33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33],
        [ 1,  2,  3,  4,  5,  6,  7,  8, 10, 11, 12, 13, 17, 19, 21, 31,  0,  2,
          3,  7, 13, 17, 19, 21, 30,  0,  1,  3,  7,  8,  9, 13, 27, 28, 32,  0,
          1,  2,  7, 12, 13,  0,  6, 10,  0,  6, 10, 16,  0,  4,  5, 16,  0,  1,
          2,  3,  0,  2, 30, 32, 33,  2, 33,  0,  4,  5,  0,  0,  3,  0,  1,  2,
          3, 33, 32, 33, 32, 33,  5,  6,  0,  1, 32, 33,  0,  1, 33, 32, 33,  0,
          1, 32, 33, 25, 27, 29, 32, 33, 25, 27, 31, 23, 24, 31, 29, 33,  2, 23,
         24, 33,  2, 31, 33, 23, 26, 32, 33,  1,  8, 32, 33,  0, 24, 25, 28, 32,
         33,  2,  8, 14, 15, 18, 20, 22, 23, 29, 30, 31, 33,  8,  9, 13, 14, 15,
         18, 19, 20, 22, 23, 26, 27, 28, 29, 30, 31, 32]])
batch with is_sorted=False = Data(x=[17, 34], edge_index=[2, 25], y=[17], train_mask=[17], input_id=[3], batch_size=3)
batch with is_sorted=True  = Data(x=[3, 34], edge_index=[2, 2], y=[3], train_mask=[3], input_id=[3], batch_size=3)

I think this edge_index is sorted by column unless I'm really mis-understanding what they mean by sorted by column. We then can see that the batch size is 17 vs 3 for is_sorted=False vs is_sorted=True. I would have expected these batch sizes to be roughly equal. Also, the number of edges in the is_sorted=false seems like its only looking at edges in the original batch.

I am running with pytorch version 1.12.1 and pytorch geometric version 2.2.0.

Any advice would be much appreciated!

[wsad1]

Your edge indices are sorted row-wise (i.e edge_index[0]) it should be sorted by column (i.e edge_index[1]). Use the coalesce with sort_by_row=False to sort by column, note: this will remove repeated edges.

[davidlow]

Thanks so much! That solved the batch size problem.

It doesn't really matter for implementation, but I'm still a little confused conceptually why this is called sort by column. Staring at the original and sorted edge_index, it looks like the sorted one has the destination node in order (assuming first row is src and 2nd row is dst) while the original has the source nodes in order. Or, the dst node is the most significant "digit" in the sorted edge_index. Is this the correct interpretation?

[wsad1]

Yeah we should probably improve the documentation to say column means destination node. We sort by destination node because neighbor loader expects it sorted by destination node, as sampling starts from destination node.