Contrastive loss DataLoader pytorch-geometric.

[nmuralid1]

Hello,
I'm quite new to pytorch-geometric and have been going through the tutorial to try and implement a data loader amenable to my current requirement. I have been looking to implement a contrastive learning solution using graphs however, in my case the positive and negative samples are known beforehand (since my problem is supervised). What is the correct (most straightforward) way to implement my pytorch geometric data loader to have a single batch in pytorch output in the following manner:

 for epoch in  epochs :
     for (batchG,batchG_Pos1,batchG_Negs1,batchG_Negs2,batchG_Negs3,...,batchG_NegsK) in batches:
             preds = contrastiveModel(batchG,batchG_Pos1,batchG_Negs1,...,batchG_NegsK)
             ####### Do something ########

I am specifically trying to implement a solution similar to Graph Contrastive Learning with Augmentations, except my augmentations are already pre-defined (and not done on-the-fly) and I want to read the dataset in using the pytorch geometric DataLoader.

In the training loop above, batchG corresponds to the batch of anchor graphs while batchG_Pos1 corresponds to the set of positive batch of graphs and each of the batchG_Negi corresponds to a set of negative graphs with the j^{th} item in each batch corresponding to the same training instance group to be used in the contrastive loss.

From what I have read so far on the Pytorch Geometric site, I am thinking of extending the Pair of Graphs example data loader to achieve this (i.e., edge_index_pos,edge_index_neg1,edge_index_neg2 ...). Is there a more straightforward (and extensible in the future) way to achieve the same effect?

[rusty1s]

I think the most straightforward way to do this is to simple hold a list of tuples, and pass this representation to the DataLoader, i.e.:

data_list = []
for i in range(num_graphs):
    data_list.append((data_i, data_i_pos_1, data_i_neg_1, ..., data_i_pos_k, data_i_neg_k))

Passing that to the DataLoader should batch these data objects along the first dimension:

loader = DataLoader(data_list, ...)

for batchG, batchG_pos_1, ..., in loader:
    ...