How to get multiple augmentations of the same graph in a single batch

[AnimeshSinha1309]

My question pertains to using Pytorch Geometric InMemoryDataset. I have a dataset of molecular graphs, and I wish to train a contrastive learning algorithm on it. I can store the original graph, and on the fly I can apply augmentations to the graph. But when feeding to my contrastive loss function, we need more than one augmentation of the same dataset element to be a part of the batch (so that both the augmentations can be pushed together in latent space).

So my question is, How do I get different on-the-fly generated augmentations of my graph (which the library calls transforms) for a single element in a single batch. Is there a clean way to do this, or do I need to hack my way through by calling the tranform function myself instead of through the library, which might well result in a loss of performance.

[rusty1s]

Mh, interesting problem. How do you solve this in regular PyTorch? One thing I could think of is to write a WrapperDataset that allows to apply different transforms to a data object:

class WrapperDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, transform1=None, transform2=None):
        self.dataset = dataset
        self.transform1 = transform1
        self.transform2 = transform2

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        data1 = self.dataset[idx]
        data2 = copy.copy(data1)
        return self.transform1(data1), self.transform2(data2)


dataset = WrapperDataset(dataset, transform1, transform2)
loader = DataLoader(dataset, ...)

for batch1, batch2 in loader:
    print(batch1)
    print(batch2)