A faster implementation of `AddMetaPaths`

[EdisonLeeeee]

🚀 The feature, motivation and pitch

I've recently come across the method AddMetaPaths and I found it quite slow particularly on large-scale heterogeneous graphs. Take an example in this discussion:

from torch_geometric.datasets import AMiner
from torch_geometric.transforms import AddMetaPaths
data = AMiner(root='./data/aminer_pyg/')[0]
metapaths = [[("author", "paper"), ("paper", "author")],
[("author", "paper"), ("paper", "venue"), ("venue", "paper"), ("paper", "author")]]
data = AddMetaPaths(metapaths)(data)

it takes me hours and much memory to run on the AMiner graph, even though I specified max_sample=1 in AddMetaPaths.

Alternatives

I think the implementation can be accelerated via multiple one-step random walks to generate the metapaths. For example, for the metapath [("author", "paper"), ("paper", "author")], we can perform one-step random walk on ("author", "paper") graph, and then on ("paper", "author") graph, to generate the metapath efficiently.

I've tested this method and it run much faster than the current implementation with lower memory overheads. Would it be possible to be added to PyG?

Additional context

No response

[rusty1s]

Definitely. Parts of the slow implementation can also be contributed to the slow CPU-based spspmm.
Adding @kpstesla to the conversation as well. Happy to take any contributions here!

[kswhitecross]

This is a great idea! I'll take a look at the implementation. I'm also curious to see how this performance compares with the much faster spspmm implementation for the gpu, although that will surely have memory limitations.

[EdisonLeeeee]

I've tested previous AddMetaPaths (implemented with spspmm) on the AMiner graph with ~300k nodes. It cannot be finished in hours on both GPU and CPU. The code snippet is:

from torch_geometric.datasets import AMiner
from torch_geometric.transforms import AddMetaPaths

data = AMiner(root="AMiner")[0]
metapaths = [
    [("author", "paper"), ("paper", "author")],
    [("author", "paper"), ("paper", "venue"), ("venue", "paper"), ("paper", "author")],
]
data = data.to('cuda')
data = AddMetaPaths(metapaths, max_sample=1)(data)