How to use transform to preserve both original data object and a copy

[KubaMichalczyk]

I have a graph for which nodes are described partially by their position in a 2-D coordinate system. What's more, the labels are invariant to symmetries. Therefore, I could easily generate more free data by reflecting nodes against x, y and xy. As a result, I have four copies of each input data object, and I want to retain all of them for training. I know the data augmentation can be done using transforms/pre-transforms. However, this usually keeps only the augmented data object, not the augmented and original objects together.
My best idea is to create a Batch object with Batch.from_data_list() to concatenate both in a super graph. This way, I can define both symmetrical reflections separately and even Compose them to generate all three reflections at one go. However, when I'm reading this in the DataLoader, the Batch object is generated, which, in fact, is composed of a batch of batches. This might be fine for some applications, but if I want to use softmax across the edges, for example, I need to know the number of edges in the original single data object, not a supergraph resulting from transformations (which can be obtained by dividing no. of edges in supergraph by batch_size). Is there another way to retain the multiple transform outputs than using a Batch object? Or maybe this is fine, and I can restore the original graph attributes (like no. of edges) in another way?

I'm attaching my implementation for horizontal reflections below (vertical reflection is analogous, and horizontal-vertical one is just a composition of the two).

class AddHorizontalSymmetry(BaseTransform):
    def __init__(self):
        super(AddHorizontalSymmetry, self).__init__()

    def forward(
        self,
        data: Union[Data, HeteroData],
    ) -> Union[Data, HeteroData]:
        data_reflected = data.clone()
        data_reflected.x[:, :, 0] = -data.x[:, :, 0]
        data = Batch.from_data_list([data, data_reflected])
        return data

[rusty1s]

How about:

dataset = torch.utils.data.ConcatDataset([
    Dataset(...),
    Dataset(..., transform=...),
    Dataset(..., transform=...),
])

[KubaMichalczyk]

this is neat, thanks!

[rusty1s]

Another option is to define a transform that randomly decides which transform to apply.