NeighborLoader may discard self-loops, and add_remaining_self_loops may not work well with trim_to_layer

[ClaudMor]

Hello,

It is my understanding that NeighborLoader samples num_neighbors nodes at each hop. But if the GNN requires self-information, it could be that during neighbor sampling the self-loops are lost (please correct me if I'm wrong on this point). The add_self_loops argument of NeighborSampler (see #64) does not seem to be available for NeighborLoader.

A possible solution to that could be to run add_remaining_self_loops after the batch has been created i.e. within the layer's forward. But this would interfere with trim_to_layer if one wishes to also use it. Infact, consider the code:

import torch_geometric as pyg
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import trim_to_layer, add_remaining_self_loops

class MyConv(MessagePassing):
    def __init__(self, other_arguments):
        # layer initialization code

     def forward(self, x, edge_index, num_sampled_nodes, num_sampled_edges, layer):
       
        trimmed_x, trimmed_edge_index, _ = trim_to_layer(layer = layer, 
        num_sampled_nodes_per_hop = num_sampled_nodes, num_sampled_edges_per_hop = num_sampled_edges, x=x, edge_index = edge_index)

        trimmed_edge_index_with_self_loops, _= add_remaining_self_loops(trimmed_edge_index)
        
        # compute embeddings and return

To compute a forward pass, the arguments would be passed to the layer's forward through the model via:

neighbor_loader = NeighborLoader( #...
batch = next(iter(neighbor_loader)
predictions = model(batch.x, batch.edge_index, batch.num_sampled_nodes, batch.num_sampled_edges)

But the call to add_remaining_self_loops would add self-loops also to nodes for which we don't want to compute the embeddings at the given layer, because trimmed_edge_index contains nodes up to the len(num_sampled_nodes) - num_layer hop in order to compute the embeddings of nodes up to len(num_sampled_nodes) -num_layer - 1 hop and add_remaining_self_loops adds self loops for every node in trimmed_edge_index, be it source or destination, as the following simple code may confirm:

edge_index = torch.tensor([[0],
                           [1]]) # 1 appears only as destination 

add_remaining_self_loops(edge_index) # tensor([[0, 0, 1], [1, 0, 1]])

Inverting the execution order of trim_to_layer and add_remaining_self_loops looks not immediate to me since the self-loops are added at the end of edge_index and this would invalidate the num_sampled_edges array (at least).

Another way I tried was to define the loader as:

import torch_geometric.transforms as T

neighbor_loader = NeighborLoader(data, 
                              num_neighbors = [ 25, 10],
                              batch_size = 512,
                              shuffle = True, 
                              input_nodes =  data.train_mask,
                              transform = T.Compose([T.AddSelfLoops(),  T.RemoveDuplicatedEdges()]))

But the num_sampled_edges of subsequent batches does not seem to keep track of the added edges.

So I'd have the following two questions:

1. Is there an easy way to efficiently add self loops in a minibatch + neighbor sampling scenario?
2. Does the problem carry over to the SparseTensor implementation of trim_to_layer (see PR)? If so, how may I get rid of it in that case?

Thank you in advance.

[rusty1s]

I think that's a great issue and something we haven't really thought to deeply about yet. One easy way to fix this would be to remove self-loops completely from your graph before inputting it into NeighborLoader, and then apply add_self_loops after each trimming stage. add_self_loops is a very fast op, and you shouldn't see any decrease in performance even if you apply it before each layer.

The better fix would be to deeply integrate it into our NeighborLoader C++ code, but we don't have any urgent plan to integrate this right now.