Using callbacks for datamodule setup preprocessing logic?

[brijow]

I have some preprocessing logic to run during datamodule setup process, but only in certain situations (mostly to try out different preprocessing steps while experimenting, but at other times, its due to the model I am using with the datamodule).

Is there a way to specify a set of data preprocessing steps to perform using callbacks? Reading the documentation, I could not find the correct hook to use.

[tchaton]

Dear @brijow,

I wondered if something like that could fit your use-case ?

class MyDataModule(LightningDataModule):

    def __init__(self):

        self._processed_train_dataset = None

    def setup(self):
        self.train_dataset = ...

    @property
    def processed_train_dataset(self):
        return self._processed_train_dataset or self.train_dataset

    @processed_train_dataset.setter
    def processed_train_dataset(self, processed_train_dataset):
        self._processed_train_dataset = processed_train_dataset
    
    def train_dataloader(self):
        return DataLoader(self.processed_train_dataset)


class Preprocessing1(Callback):

    def preprocess_function(self, dataset):
        # your preprocessing logic

    def on_train_start(self, trainer, pl_module):
       # apply processing
        trainer.datamodule.processed_train_dataset = self.preprocess_function(trainer.datamodule.train_dataset)
        
        # force dataloader reload
        # trainer.reset_train_dataloader(pl_module)


trainer = Trainer(callbacks=Preprocessing1())
trainer.fit(model, dm)

[brijow]

Thanks @tchaton . The code you suggested got me on the right track -- I ended up using the setup hook. Here's what I ended up using, however I haven't tried it in any distributed modes, so I'm not sure its bullet proof but for now seems to work well enough.

Just to give a bit of context, I'm using a HuggingFace DatasetDict, so I'm applying the preprocessing logic to all my splits by simply calling map() on the DatasetDict .

from pytorch_lightning.callbacks import Callback

class PreprocessData(Callback):

    def setup(self, trainer, pl_module, stage=None):
        if stage != "fit" and trainer.datamodule._has_setup_fit:
            return

        def preprocess_func(batch):
            # preprocess logic here ....
            return batch

        trainer.datamodule.dataset = trainer.datamodule.dataset.map(
            preprocess_func, batched=True
        )

[tchaton]

Yes, I think it is fine. Each ranks will have its own copy of the processed data, and then it should be wrapped into a DistributedSampler automatically :)