SparseTensor support in HeteroData

[Vid201]

Hi, I am trying to use HeteroData object on MAG240M from OGB-LSC and I am having some troubles (not sure what I am doing wrong). The first question would be if SparseTensor for edge_index (or adj_t maybe) in HeteroData is supported? I did not found any example in tests or docs.

[rusty1s]

Yes, this is supported. Each occurrence of edge_index can simply be replaced by a SparseTensor variant adj_t. This way, you can create the heterogeneous graph in a similar fashion as we do in our OGB mag240m example. Keep in mind that loading node features into memory will consume a lot of RAM.

[Vid201]

Thanks! ToSparseTensor is what I was looking for.

Does it make sense that performance is now affected:

1. no conversion and pass batch.edge_index_dict to model -> 5 it/s
2. convert to SparseTensor and pass batch.edge_index_dict (torch tensor) to model -> 2.5 it/s
3. convert to SparseTensor and pass batch.adj_t_dict (SparseTensor) to model -> 1.75 it/s

[rusty1s]

This shouldn't be the case :) To which example are you referring to?

[Vid201]

https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/to_hetero_mag.py - using NeighborLoader.

[rusty1s]

Let me check and come back to you :)

[Vid201]

Sure, thank you very much :)

[rusty1s]

For me, using ToSparseTensor() with batch.adj_t_dict is equally fast. I think for NeighborLoader there shouldn't be a drastic difference in runtime, as you cannot really exploit the memory-efficiency of SparseTensor in local subgraphs where num_nodes ~ num_edges.

It is a bit weird that using SparseTensor is so much slower for you. Are you running on CPU (since I'm at 20-25 it/s)? Note that SparseTensor and passing edges via batch.edge_index_dict is not recommended, as it will convert the SparseTensor representation back to the edge_index one.

Edit: On CPU, using SparseTensor is twice as fast for me. Mmh.

[Vid201]

Thanks for testing. I tried on the different machine with less overall CPU usage (some other things are also running on the first machine) and now the performance is the same in both cases with much higher it/s. I guess copying the SparseTensor from CPU to GPU is more CPU intensive? Don't know.

In all cases, the models were running on GPU.

Thanks again for the help!