LinkNeighborLoader not aware of negative samples from RandomLinkSplit

[emilioangelina]

I want to train a GAE model for generating embeddings from a networks. In order to fit the network in the GPU memory I need to split edges in tran / val / test and then generate mini-batches of links.
I have followed the general procedure explained in several posts for spliting links with RandomLinkSplit, using a negative sampling ratio of 1 for validation and testing data. By inspecting the edge_label tensor I can see that it contains 1s and 0s, as expected.
Then I pass the splits to LinkNeighborLoader, this time I set neg_sampling_ratio=0 as I already generated the negative sampling with RandomLinkSplit. However, the generated batches only have the positive samples. When I set the neg_sampling_ratio=1 in LinkNeighborLoader, the negative samples seems to be added as the size of edge_label and edge_label_index tensors are doubled in size. But the edge_labels tensor is populated of 2s and 0s instead of 1s and 0s like in the original tensor generated by RandomLinkSplit. I am pasting below the relevant code:

from torch_geometric.data import Data
data_gae = Data(x=identity_matrix, edge_index=edge_index,edge_attr=edge_weights)
from torch_geometric.transforms import RandomLinkSplit
from torch_geometric.loader import LinkNeighborLoader

Split the links into tran, val and test

transform = RandomLinkSplit(num_val=0.1, num_test=0.1, neg_sampling_ratio=1.0, disjoint_train_ratio=0.3, add_negative_train_samples=False, is_undirected=True)
train_data, val_data, test_data = transform(data_gae)

Batch links so that they can fit into the GPU memory

batch_size = 128
num_neighbors = [-1]

train_loader = LinkNeighborLoader(train_data, num_neighbors=num_neighbors,
batch_size=batch_size,
neg_sampling_ratio=1,
edge_label_index=train_data.edge_label_index,
edge_label=train_data.edge_label,
shuffle=True)

val_loader = LinkNeighborLoader(data=val_data,
num_neighbors=num_neighbors,
neg_sampling_ratio=0,
edge_label_index=val_data.edge_label_index,
edge_label=val_data.edge_label,
batch_size=batch_size,
shuffle=False)

test_loader = LinkNeighborLoader(test_data,
num_neighbors=num_neighbors,
neg_sampling_ratio=0,
edge_label_index=test_data.edge_label_index,
edge_label=test_data.edge_label,
batch_size=batch_size, shuffle=False)

Thanks for your help

[rusty1s]

But the edge_labels tensor is populated of 2s and 0s instead of 1s and 0s like in the original tensor generated by RandomLinkSplit.

Yeah, this is expected, since LinkNeighborLoader just increments any existing edge_label and represents its own negatives as 0.

If you are using LinkNeighborLoader(neg_sampling_ratio=0), you should definitely see edge_label containing both zeros and ones.

[emilioangelina]

Thank you. Because I had set shuffle to False, the batches I inspected exclusively consisted of 1s, with the 0s appearing at the end of the validation data. However, upon setting neg_sampling_ratio to 0 and shuffle to True, I discovered that indeed the batches contained both 1s and 0s.