Question concerning the function "train_test_split_edges"

[bits-glitch]

Hi PyG team,

I was looking at the documentation of the function torch_geometric.utils.train_test_split_edges.

Can you tell me, why this function computes:

val_pos_edge_index, val_neg_edge_index
test_pos_edge_index,test_neg_edge_index

but

train_pos_edge_index, train_pos_neg_adj_mask

What is the reason for computing the adjacency mask instead of the edge_index format?

Thank you very much for the help!

[rusty1s]

The train_pos_neg_adj_mask can be used to sample negative edges. If you want to convert that into an edge_index representation, you can do the following:

train_neg_edge_index = data.train_neg_adj_mask.nonzero(as_tuple=False).t()

However, note that we typically cannot train against all negative edges since the amount of negative edges is quite large in general.

[bits-glitch]

Hey! Thank you very much for the answer, that clarifies my question :) !

From what I understand, the train_pos_neg_adj_mask seems a more memory efficient solution, thus I would like to keep it that way in my dataset, which consists of multiple graphs. However, (I want to store each graph with its corresponding splits in a data_list), I get dimensionality problems with the collate function when I want to store the data.train_neg_adj_mask of each graph in the same data_list (list of Data objects). Can you recommend a solution to this problem?

[rusty1s]

The only way possible to batch dense adjacency matrices is by padding them to equal size :(
In this case, I think it's best to drop train_neg_adj_mask completely, and use torch_geometric.utils.batched_negative_sampling to create negative samples directly during training. Does that work for you?

[bits-glitch]

Yeah that works for me! Thank you very much for your time and the recommendation to use torch_geometric.utils.batched_negative_sampling!