[question] How to override how `self.train_transform` is applied to batches in `LightningDataModule`.

[akashpalrecha]

❓ Questions and Help

What is your question?

I want to apply the same random transformation on both X and Y (both are batches if RGB image) of a training batch using Kornia after the batch has been moved to the GPU (or TPU).

I cannot figure out what part of PL I need to override to make this work.

As far as Kornia is concerned, I know how to get the state of a random transform and use it again for another batch. But I cannot understand where inside LightningDataModule I am supposed to write that code such that it works properly even for MultiGPU training.

Code

What have you tried?

I can simply include the augmentation step in the LightningModule's training step after receiving the batches, but that would decouple my LightningDataModule from it's transforms which I don't want to do.

Any help would be greatly appreciated !! 🙂

On another note, the documentation isn't very clear on how the self.train_transform property of LightningDataModule is applied to batches of data and how we can change that behavior (which is pretty much what my question is).

What's your environment?

• OS: [Ubuntu 16.04 LTS]
• Packaging [ conda ]
• Version [0.8.5]

[github-actions]

Hi! thanks for your contribution!, great first issue!

[rohitgr7]

you can access that using self.trainer.datamodule.train_transforms within LightningModule but I am not sure if it's the right way to do it.

[akashpalrecha]

@rohitgr7 This is what I'm having to do right now. I'm calling the setup function on my LightningDataModule first, and then setting a transforms property in the LightningModule which I get from my LightningDataModule. I use the transforms to then manually transform my image pairs in every training step.

But ideally, I'd want the LightningDataModule to take care of it internally.

[rohitgr7]

@nateraw

[versatran01]

I use kornia and I just wrote my own transform and call it in train_step.
Kornia's design goal is to make all these transforms differentiable, so it's not that bad to call it in the train_step.

[akashpalrecha]

@versatran01 Yeah, that's what I'm doing right now and there's really no issue with this right now to be honest. But it prevents me from using the same data module design in other lightning modules without having to manually add the transformation code in training_step every time. Ideally, the data module should do everything related to data.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!