NeighborLoader for multiple graphs

[ilyalasy]

So, I know that NeighborLoader mostly needed for large scale graphs, but it should work as well for collection of multiple graphs, right?

I have a 3 layer GATConv model just like here. I also have a dataset of ~900 graph objects with labeled nodes. If I use simple DataLoader, model trains successfully. If I switch to NeighborLoader loss is not changing much.

I input all graph dataset into NeighborLoader, while filtering nodes by some criteria (idea is to undersample majority class that way):

train_obj = Batch.from_data_list(train_dataset) # loading all graphs into batch object
indices = torch.where(train_obj.y[:, 1]==1)[0] #filtering majority class
train_loader = NeighborLoader(train_obj,num_neighbors=[10]*3, input_nodes=indices, batch_size=1024, shuffle=True)  

Usual training code:

for data in train_loader: 
  self.optimizer.zero_grad() 
  data = data.to(device)
  batch_size = data.batch_size
  
  y_cls = data.y[:, 0]
  
  out_cls = self.predictor(data.x, data.edge_index, data.batch)[:batch_size]
  
  loss = self.criterion(out_cls, y_cls[:batch_size])
  loss.backward()
  self.optimizer.step()

Any advice on this situation? Is it correct to use NeighborLoader in such a way?
Thank you!

[wsad1]

The usage of NeighborLoader looks right to me. Maybe something else is causing this problem, did you use the usual DataLoader also while removing the majority class.

[ilyalasy]

No, while using DataLoader i can't remove majority class as i'm loading whole graphs so all nodes were inputted to the model. You think issue can be in lack of those nodes i remove... Probably, but they all belong to majority class, and still after filtering there plenty of such nodes, so its strange

[rusty1s]

Is it possible for you to check what's the performance of a pre-trained model (trained in full-batch mode via DataLoader) when inputted into NeighborLoader? At best, once with sampling all edges [-1] * 3, and once with sampling only [10] * 3 edges. When sampling all edges , performance should be identical. If the performance drops when only sampling a fraction of neighbors, this might indicate that full neighborhood information is really important.

[ilyalasy]

I'll try that, thank you for idea!