NeighborLoader for heterogeouse graph link prediction task: how to update edge_label_index

[zqfang]

I got a huge heterogouse graph for link predicition. After I split the graph using RandomLinkSplit, I used NeighborLoader or HGTLoader to subgraph for training on GPU.

Everything works fine, but only the edge_label_index created by RandomLinkSplit is not updated when call NeighborLoader, causing training process stopped.

Did I do something wrong ?

example code:

torch.manual_seed(seed=123456)
train_data, val_data, test_data = T.RandomLinkSplit(is_undirected=True, 
                                         add_negative_train_samples=True, 
                                         neg_sampling_ratio=1.0,
                                         key = 'edge_label',
                                         rev_edge_types=('mesh','rev_genemesh','gene'))(gm_data)

train_loader = NeighborLoader(train_data, 
                              num_neighbors = {key: [15] * 3 for key in train_data.edge_types},
                              shuffle=True, batch_size=128,
                              input_nodes=('mesh',None), )

# training
for epoch in range(1, 10):
    for i, batch in enumerate(train_loader):
        optimizer.zero_grad()    
        batch.to(device)
        out_dict = model(batch.x_dict, batch.edge_index_dict)
        batch_size = batch['mesh'].batch_size
        ## FIXME: this won't work because the edge_label_index_dict never updated
        loss = model.score_loss(out_dict, batch.edge_label_index_dict, batch.edge_label_dict)
        loss.backward()
        optimizer.step()
        train_loss = loss.item()
        #val_loss, aupr, auroc = valid(val_data, device)  
        print(f"Batch {i:03d}, Train Loss: {train_loss:7.f}")

[rusty1s]

Yeah, that's a current limitation as we can only filter out vectors that correspond to the shape of nodes or edges. We have plans to integrate a NeighborLoader for link prediction tasks as well. One way around this is to inhereit the NeighborLoader by yourself, similar to what we are doing in the OGB link prediction example, see here

[zqfang]

Great. Thank you very much. Hope the function will be implemented soon.

By the way. is there a way to suport multiple node_types input for NeigborLoader(input_nodes= ....). This will be useful when sampling a bipartile graph @rusty1s

[rusty1s]

Can you elaborate on the use-cases when this might be useful? In general, when doing a node classification task, we only want to obtain output embeddings for a single node type. I agree that this might be useful to have in link prediction tasks though.