Workaround for link-based data loader on heterogeneous graph

[ayushnoori]

Hi PyG team,

Thank you very much to @rusty1s and team for your work in developing an incredible PyG resource and community.

I'm attempting to design an unsupervised link prediction training task. My data, a heterogeneous graph with multiple node and edge types, is stored as a HeteroData object. Following the NeighborLoader examples in the documentation here and this tutorial, I attempted to use a LinkNeighborLoader data loader for this task.

However, I realized that only single node/edge types can be used as seed nodes in LinkNeighborLoader and related data loaders (see discussions at #7551 and #5269, and the open PR at #5521).

Do you have any suggestions for a workaround data loader that can operate on a heterogeneous graph for link prediction?

In particular, I am attempting to use a Heterogeneous Graph Transformer supported by the HGTConv class (as documented here and with the hgt_dblp.py demo here). I had explored using HGTLoader, but don't see negative sampling supported, which I would need for link prediction. I would appreciate any recommendations and/or pointers to resources where I can learn more.

Thanks for your guidance!

P.S., for context, I am currently constructing the HeteroData object and splitting my graph into train/validation/test using the approach outlined in this Gist:

https://gist.github.com/ayushnoori/1c2430f6e489f05631c4f304a614fbc2

Please let me know if my approach is incorrect. Thanks!

[rusty1s]

This is sadly not yet possible, sorry. The only workaround would be to map your HeteroData into a "typed" homogeneous graph by called data.to_homogeneous(). Would that workaround work for you?

[ayushnoori]

Hi Matthias, thank you for letting me know. Unfortunately, I think data.to_homogeneous() will not work because I'd like the result of each batch created by iterating through the DataLoader to be a HeteroData object, so I can take advantage of the HGTConv operator (as in hgt_dblp.py).

I also noticed that, by default, data.to_homogeneous() does not preserve the node and edge feature matrices? That is,

data_hg = data.to_homogeneous()
data_hg

This yields the following, which lacks the .x and .edge_attr attributes.

Data(edge_index=[2, 4460890], edge_label=[3186322], node_type=[132662], edge_type=[4460890])

I guess I would have to reconstruct them manually?

I had also thought about constructing a separate LinkNeighborLoader object for each edge type, then sequentially iterating through them as described here. This is not ideal, because I was seeking a single DataLoader object (e.g., for use with PyTorch Lightning) and also because it might be better to shuffle edge type across seed edges. What do you think about this approach?

Do you have an estimated timeline for when this feature might be fully supported by PyTorch Geometric?

Thanks again for your help and advice.

[rusty1s]

I guess I would have to reconstruct them manually?

In PyG master, these attributes should now be included by default and padded in case the feature size is not the same.

Constructing separate LinkNeighborLoader classes is possible, but a very hacky workaround :(

I don't have an estimated timeline, sadly. pyg-lib supports this already out of the box, but there are some changes we need to make on the LinkNeighborLoader side to support this fully.

[ayushnoori]

Thanks for letting me know, Matthias. Let me know if I can help with or support any changes to LinkNeighborLoader. In the meantime, will keep this thread open. Happy to also convert to an issue if you'd like for tracking.

[rusty1s]

It is probably a good idea to start with this effort for NeighborLoader rather than LinkNeighborLoader since in general this is easier to handle. In either case, the input nodes/edges needs now to be given as a dictionary of node types/edge types and their corresponding input node/edge values, and this change would need to be propagated to NodeSamplerInput and EdgeSamplerInput (see here).

[Mrugankakarte]

Hi,

Are there any official examples for this feature if it is available in recent versions?