LinkNeighborLoader doesn't sample edges more than 2 hops though I set it.

[Xurorx]

Hello,

I would like to implement an RGCN model for a link prediction task on heterogenous graph OGB-MAG dataset.
I use RandomLinkSplit to split the graph into train, val and test. Then I use LinkNeighborLoader to make train_loader, val_loader and test_loader. However, LinkNeighborLoader doesn't work as I expected.

import torch
import torch_geometric
from torch_geometric.datasets import OGB_MAG, Planetoid
from torch_geometric.transforms import RandomLinkSplit
from torch_geometric.transforms import ToUndirected
from torch_geometric.loader import LinkNeighborLoader

print(f"torch version: {torch.__version__}")
print(f"torch geometric version: {torch_geometric.__version__}")

I have GPU resources, just use CPU on my local machine to start.

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

dataset = OGB_MAG(root='./data', preprocess='metapath2vec')
data = dataset[0]

data

HeteroData(
paper={
x=[736389, 128],
year=[736389],
y=[736389],
train_mask=[736389],
val_mask=[736389],
test_mask=[736389]
},
author={ x=[1134649, 128] },
institution={ x=[8740, 128] },
field_of_study={ x=[59965, 128] },
(author, affiliated_with, institution)={ edge_index=[2, 1043998] },
(author, writes, paper)={ edge_index=[2, 7145660] },
(paper, cites, paper)={ edge_index=[2, 5416271] },
(paper, has_topic, field_of_study)={ edge_index=[2, 7505078] }
)

train_data, val_data, test_data = transform(data)

train_data

HeteroData(
paper={
x=[736389, 128],
year=[736389],
y=[736389],
train_mask=[736389],
val_mask=[736389],
test_mask=[736389]
},
author={ x=[1134649, 128] },
institution={ x=[8740, 128] },
field_of_study={ x=[59965, 128] },
(author, affiliated_with, institution)={
edge_index=[2, 584640],
edge_label=[501120],
edge_label_index=[2, 501120]
},
(author, writes, paper)={
edge_index=[2, 4001570],
edge_label=[3429916],
edge_label_index=[2, 3429916]
},
(paper, cites, paper)={
edge_index=[2, 3033112],
edge_label=[2599810],
edge_label_index=[2, 2599810]
},
(paper, has_topic, field_of_study)={
edge_index=[2, 4202845],
edge_label=[3602438],
edge_label_index=[2, 3602438]
}
)

# Define seed edges:
edge_label_index = train_data["author", "writes", "paper"].edge_label_index
edge_label = train_data["author", "writes", "paper"].edge_label

train_loader = LinkNeighborLoader(
    data=train_data,
    num_neighbors=[5,10,15],
    edge_label_index=(("author", "writes", "paper"), edge_label_index),
    edge_label=edge_label,
    batch_size=128,
    directed=True,
    shuffle=True,
)

# Inspect a sample:
sampled_data = next(iter(train_loader))

print("Sampled mini-batch:")
print("===================")
print(sampled_data)

Sampled mini-batch:

HeteroData(
paper={
x=[4036, 128],
year=[4036],
y=[4036],
train_mask=[4036],
val_mask=[4036],
test_mask=[4036]
},
author={ x=[4864, 128] },
institution={ x=[0, 128] },
field_of_study={ x=[0, 128] },
(author, affiliated_with, institution)={
edge_index=[2, 0],
edge_label=[501120],
edge_label_index=[2, 501120]
},
(author, writes, paper)={
edge_index=[2, 5511],
edge_label=[128],
edge_label_index=[2, 128],
input_id=[128]
},
(paper, cites, paper)={
edge_index=[2, 4960],
edge_label=[2599810],
edge_label_index=[2, 2599810]
},
(paper, has_topic, field_of_study)={
edge_index=[2, 0],
edge_label=[3602438],
edge_label_index=[2, 3602438]
}
)

As you can see, a batch from train_loader only contains "(author, writes, paper)" and "(paper, cites, paper)" edges. I thought LinkNeighborLoader will samples multiple hops from both ends of a link. For example, it should sample "author, writes, paper" -> "paper, cites, paper" -> "paper, has, topic". I also want it to have edges "author, affiliated with, institution". But it doesn't sample those edges, is there anything wrong with my code?

Thank you in advance!

[rusty1s]

You may want to create bi-directional edge types for all the edges such that all node types are reachable. You can do this via the T.ToUndirected transform.

[Xurorx]

Thank you @rusty1s !