NeighborLoader seed nodes not the first nodes in the batch

[Romain-Nicolle]

Hi everyone,

I have a problem with the NeighborLoader() and I need some assistance.
I am using an Heterogeneous Graph with 5 node_types and 20 edge_types to perform node classification.

I am using these versions of these packages:

torch.__version__ = 2.1.2+cu118
torch_geometric.__version__ = 2.4.0
pyg_lib.__version__ = 0.3.1+pt21cu118

My HeteroData looks like that:

HeteroData(
  NODE_TYPE_A={
    x=[30391, 8],
    y=[30391],
	num_nodes=30391,  
	train_mask=[30391],
	test_mask=[30391],
	global_mask=[30391],
  },
  NODE_TYPE_B={
    num_nodes=14337,
    x=[14337, 687],
  },
  NODE_TYPE_C={
    num_nodes=27993,
    x=[27993, 199],
  },
  NODE_TYPE_D={
    num_nodes=4039,
    x=[4039, 425],
  },
  NODE_TYPE_E={
    num_nodes=11271,
    x=[11271, 349],
  },
  (NODE_TYPE_A, edge_type_1, NODE_TYPE_B)={ edge_index=[2, 3000000] },
  ...
  (NODE_TYPE_E, edge_type_20, NODE_TYPE_A)={ edge_index=[2, 15000] }
 )

Due to the size of my HeteroData object, I need to perform mini-batch during training so I am using different Neighborloaders, more specifically I am using 3:

• one train_loader for training. I use a sampling strategy there because my training set is imbalanced. In my experiment, the first ~20000 nodes are the training nodes and the last ~10000 are the testing nodes. So my train_mask is of size ~30000 with 20000 times 1, followed by 10000 times 0. I will change the train/test split later on to perform cross_validation.
• one test_loader for testing. My test_mask is of size ~30000 with 20000 times 0, followed by 10000 times 1.
• one global_loader to get the final output tensor of every node in the graph in order to save it in a dataframe. My test_mask is of size ~30000 with 30000 times 1. global_mask = train_mask + test_mask

I will be using the HeteroConv layer and I want to get message_passing from the direct neighbors as well as the neighbors of the neighbors (2-hops away per edge_type). So my code looks like that:

global_loader = NeighborLoader(data, 
                        num_neighbors = {key: [25] * 2 for key in data.edge_types},
                        batch_size = 4,
                        input_nodes=('NODE_TYPE_A', data['NODE_TYPE_A'].global_mask),
                        num_workers = 6,
                        filter_per_worker = None
                        )
sampler = ImbalancedSampler(data['NODE_TYPE_A'].y[:20000].long())

train_loader = NeighborLoader(data, 
                        num_neighbors = {key: [25] * 2 for key in data.edge_types},
                        batch_size = 32,
                        input_nodes=('NODE_TYPE_A', data['NODE_TYPE_A'].train_mask),
                        num_workers = 6,
                        sampler = sampler,
                        filter_per_worker = None
                        )

test_loader = NeighborLoader(data, 
                        num_neighbors = {key: [25] * 2 for key in data.edge_types},
                        batch_size = 32,
                        input_nodes=('NODE_TYPE_A', data['NODE_TYPE_A'].test_mask),
                        num_workers = 6,
                        filter_per_worker = None
                        )

It seems to me, according to the documentation of the NeighborLoader there, that the seed_nodes (the input_id) are supposed to be the first nodes in the batch.
Unfortunately, that is not always the case in my examples.

I mean, for some of the loaders it is true but not for all. It is causing errors in my code and I don't know what to do.

batches = [i for i in global_loader]
i = 0
batch_size = batches[i]['GENE'].batch_size
print(batches[i]['GENE'].input_id)
print(batches[i]['GENE'].n_id[:batch_size+8])

Will have this output:

tensor([0, 1, 2, 3])
tensor([   1,    1,    1,    1,  824, 2128, 2934, 3067, 4248, 5387, 7737, 9147])

So the input_id are not the first nodes in the batch ? And these lines:

i = 1
batch_size = batches[i]['GENE'].batch_size
print(batches[i]['GENE'].input_id)
print(batches[i]['GENE'].n_id[:batch_size+8])

Will have this output:

tensor([4, 5, 6, 7])
tensor([   1,    1,    1,    1,  824, 2128, 2934, 3067, 4248, 5387, 7737, 9147])

So the first_nodes in the batch are the same as in the previous batch ?

But for the train_loader:

batches = [i for i in train_loader]
i = 0
batch_size = batches[i]['GENE'].batch_size
print(batches[i]['GENE'].input_id)
print(batches[i]['GENE'].n_id[:batch_size+8])

Will have this output.

tensor([  944,  7408,  5559,  1991, 15772,  6014, 16835,  2346,  1634, 13081,
          797, 18534,  3873,  9641,  9882,  5995,  4196,  6096,  5995,  7803,
        10572, 16717,  3766,    85, 18041,  4086, 16423,  2526, 10483,  1261,
         1245, 15909])
tensor([  944,  7408,  5559,  1991, 15772,  6014, 16835,  2346,  1634, 13081,
          797, 18534,  3873,  9641,  9882,  5995,  4196,  6096,  5995,  7803,
        10572, 16717,  3766,    85, 18041,  4086, 16423,  2526, 10483,  1261,
         1245, 15909,  8123,  8993,  3406,  1851,  1473,  9697,  8860, 10157])

And

i = 1
batch_size = batches[i]['GENE'].batch_size
print(batches[i]['GENE'].input_id)
print(batches[i]['GENE'].n_id[:batch_size+8])

Will have this one

tensor([ 9919,  4115,  7419, 13957, 11574,   305, 10016, 13461,  7537,  1787,
        11638,  8741,   480,  6320, 12802,  5416,  5206,  7191, 10574,   397,
         2379,  6210,  2997, 11388, 11024,  1900, 16058,   464,  4405,  6346,
        14008,  9087])
tensor([ 9919,  4115,  7419, 13957, 11574,   305, 10016, 13461,  7537,  1787,
        11638,  8741,   480,  6320, 12802,  5416,  5206,  7191, 10574,   397,
         2379,  6210,  2997, 11388, 11024,  1900, 16058,   464,  4405,  6346,
        14008,  9087,    70,   512,  1086,  1202,  1657,  2878,  3132,  3203])

So the train_loader is working.
But the global_loader and the test_loader are not behaving properly.

Has anyone already experienced that ?
Do you know what I should do to correct that ?
A combinations of parameters maybe ? num_workers, shuffle, replace etc ?

Any help would be greatly appreciated.

Thanks

[Romain-Nicolle]

Update: to test whether or not it was a problem with my data object/code, I tried to rerun this script provided as an example for the NeighborLoader in Heterogeneous Graph.

I run the beginning of the script and these lines:

train_input_nodes = ('paper', data['paper'].train_mask)
val_input_nodes = ('paper', data['paper'].val_mask)
kwargs = {'batch_size': 1024, 'num_workers': 6, 'persistent_workers': True}

train_loader = NeighborLoader(data, num_neighbors=[10] * 2, shuffle=True,
                                  input_nodes=train_input_nodes, **kwargs)
batch = next(iter(train_loader))
batch_size = batch['paper'].batch_size

tmp = (batch['paper'].input_id == batch['paper'].n_id[:batch_size])
print(f"This tensor is of shape: {tmp.shape}")
print(f"The dtype is: {tmp.dtype}")
print(f"The sum is: {tmp.sum()}")

And the output is:

This tensor is of shape: torch.Size([1024])
The dtype is: torch.bool
The sum is: 0

This seems to confirm, that the seed_nodes (i.e. those given by input_id) are not the first nodes in the batch in my instance.
Did I misunderstand what it should be ? Or do you think that there is a problem with my pyg installlation ? Or pyg-lib, torch-scatter etc. ?

Thank for any help.

Best,

EDIT: typo