How to minibatch heterogeneous networks for link prediction?

[saiden89]

Hi everyone,
I have a undirected heterogeneous graph with 1 type of node, but different kinds of edges, with the ultimate goal of performing link regression (as it is a probabilistic network with link weights between 0 and 1) on all edge types.

The current workflow first applies a RandomLinkSplit transformation to get the train, validation and test splits and then instantiates a NeighborLoader on each of the splits to get their respective loaders, which are then passed to a Pytorch Lightning model.
I can change the node loaders' parameters just fine in order to get different batch_sizes, but would it be possible to batch the edges instead? I think this would pose some problems since the different edge types can differ for the number of elements even by orders of magnitude. Even so, should I then create different LinkLoaders for each edge type?
I'm not even sure that batching for edges even makes sense, and I would greatly appreciate your input on the matter!

[rusty1s]

I think what you want to do is use LinkNeighborLoader instead of NeighborLoader for this.

[saiden89]

Would you know how to retrieve the edge_weight for the different edges? If I happen to use a different decoder for each edge type I would need to know the edges-decoder mapping.

[rusty1s]

Wouldn't you still be able to use edge_weight and edge_weight_index since they are basically untouched?

[saiden89]

For the train set yes, the transformation is reversible and produces the same heterogeneous object back. But for the validation and test sets, only the edge_index gets mapped back to each specific edge_type. Indeed all the edge_weight get merged together with no index to map the back to the edge types they belong to. Besides the impossibility to use edge-specific decoders, it is also impossible to produce metrics on a per-edge basis.
I'm not sure what you refer to with "edge_weight and edge_weight_index are basically untouched", but maybe I am missing the point completely.
Thank you for your help!

[rusty1s]

The problem with this approach is that the edge_type mapping is lost for validation and test since they have a different shape. Why don't you just do:

data = FakeHeteroDataset(num_graphs=1, num_node_types=1, num_edge_types=4,
                         edge_dim=1, node_label=0)

transform = RandomLinkSplit(
    key="edge_weight",
    add_negative_train_samples=False,
    neg_sampling_ratio=0,
    edge_types=data[0].edge_types,
)

train, val, test = transform(data[0])
print(val)

[saiden89]

Because the original point of converting into a "batched" homogeneous network was to enable the creation of one dataloader per train, validation, and test split which include all the edge types, which is then passed to the Lightning model. To my understanding, since the model has different decoders and (possibly) decoders per edge type, it should be able to reconstruct the "original" heterogeneous graph before training. This currently works for the training set but not for the validation and test sets.
I am summarizing the features here as they are spread across several comments:

• Goal: implement edge mini batching
• Data is heterogeneous with 4 edge types and 1 node type
• The objective is edge regression. Edges have associated probabilistic weights between 0 and 1.
• Model is a Lightning model because of distributed training
• Model has edge-specific encoders, an aggregator, and edge-specific decoders

If helpful I can share my progress so far with the model. I would also be happy to contribute with an example to help people with similar settings once the matter is settled.
Thank you again for your time and I apologize if my problem was unclear from the beginning.