Accessing DataModule or DataLoaders within model hooks

[MohammedAljahdali]

Hey, as the title says, I want to access the DataModule or the DataLoader inside the on fit start hook. Is this possible and how can I do it? To be more specifc I want to access my model, when I have access to my DataModule, to get a batch of data, then use it to apply some pruning algorithm on my model.

[awaelchli]

self.datamodule or self.trainer.train_dataloader in the LightningModule

[MohammedAljahdali]

So, after testing what you suggested I got the following results:

def on_fit_start(self):
        print(self.datamodule)  # None
        print(self.trainer.train_dataloader) # None
        print(self.trainer.datamodule) # <src.datamodules.classification_datamodule.ClassificationDataModule object at 0x7fba50037ca0>

[awaelchli]

Yep I guess on_fit_start is too early then.
Needs to be on_train_start() 😅

[MohammedAljahdali]

Do you recommend I move my pruning logic to on_train_start(), or keep it at on_fit_start() and use self.trainer.datamodule is there any advantages or disadvantages? Thank you for your help!

[tchaton]

Hey @MohammedAljahdali,

self.trainer.train_dataloader is actually a CombinedDataLoader wrapping the datamodule.trainer_dataloader(). I guess it depends what you want to do with your Callback exactly :)

You could also use the on_train_batch_start hook which has access to the current batch.

    def on_train_batch_start(
        self,
        trainer: 'pl.Trainer',
        pl_module: 'pl.LightningModule',
        batch: Any,
        batch_idx: int,
        dataloader_idx: int,
    ) -> None:
        """Called when the train batch begins."""
        pass

quick note: Lightning provides a ModelPruning Callback to prune your model while training: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/pruning.py#L57

[MohammedAljahdali]

Thank you @tchaton,

I assume that on_train_batch_start() gets called at the start of every training step right? What I want to do is at the start of the pruning run, I want to execute some pruning algorithm, then continue training the model.

I have taken a look at ModelPruning, but currently, I am moving a pruning project into Lightning, all the pruning logic is already available, and all I need to do is integrate things.