Datamodule and multiple transforms

[abcamiletto]

Hi all,
I have a Dataset which needs a different augmentation for the training and the validation set. I cannot find a straightforward way to do that. My implementation is pretty much the following

class MyDataModule(pl.LightningDataModule):
    def __init__(self, train_transform, val_transform):
        self.dataset = MyDataset(train_transform, val_transform)
        self.train_set, self.val_set, self.test_set = torch.utils.data.random_split(self.dataset, [0.8,0.1,0.1])

    def train_dataloader(self):
        return DataLoader(self.train_set)

    def val_dataloader(self):
        return DataLoader(self.val_set)

    def test_dataloader(self):
        return DataLoader(self.test_set)

I need to call self.dataset.use_validation_transform() at the beginning of each validation (at the end of each epoch).
I expected to be able to do so in the stage method but found out that the stage method it's only called at the beginning of the .fit() call.

I'm aware that i could call that from the model accessing it by the self.datamodule attribute in the on_validation_start method, but it seems an antipattern to me. I'd love to keep model and data independent.

Is there an easy way to do so that i'm missing? Or am i doing something weird enough that I should avoid?
Thx in advance for any response!

[akihironitta]

You can simply replace the original transform with a preferred transform for each dataset instance after splitting, no?
e.g. self.val_set.transform = val_transform

Or, you can do something similar to an example in Bolts which instantiates dataset multiple times to create datasets with different transforms for each with the seed fixed: https://github.com/Lightning-AI/lightning-bolts/blob/ad771c615284816ecadad11f3172459afdef28e3/pl_bolts/datamodules/stl10_datamodule.py

[abcamiletto]

First of all, thanks a lot for the response!

You can simply replace the original transform with a preferred transform for each dataset instance after splitting, no? e.g. self.val_set.transform = val_transform

self.val_set is a Subset object, and it shares its Dataset object with train_set and test_set. So i don't think thats a viable option: if i change self.val_set.dataset.transform i will end up modifying it for every set

Or, you can do something similar to an example in Bolts which instantiates dataset multiple times to create datasets with different transforms for each with the seed fixed: https://github.com/Lightning-AI/lightning-bolts/blob/ad771c615284816ecadad11f3172459afdef28e3/pl_bolts/datamodules/stl10_datamodule.py

That is more or less how i actually got it working, but i'm looking forward to load all the dataset in memory and I would prefer not to have any duplicate.

I guess my best chance is splitting the dataset before hand and create 3 separate Dataset objects

[akihironitta]

I guess my best chance is splitting the dataset before hand and create 3 separate Dataset objects

I'd suggest the same 👍 (hoping that someone has a better idea)