Mechanism to skip certain hooks

[awaelchli]

🚀 Feature

Do we need a way to prevent certain hooks to be executed? I'm not entirely sure how solid this idea is, so I'm hoping for some discussion :)

Motivation

A user encountered a use case in which they wanted to build the model in the setup hook. However, because the setup hook is exectued everytime regardless whether the model was already trained or not, this would then overwrite the weights of the model, making continued training impossible. #5410

Pitch

Something abstract like
model.skip_hooks = ["setup", ...]
could be considered.

Alternatives

The user can handle it on their side with some conditional code in their hook,

def setup(self, stage):
    if self.has_setup:
        return
    # do your thing
    ...
    self.has_setup = True

or forcefully remove the method from the model object, delattr(model, "setup"), for example.

Additional context

Originates from here

cc: @sjgosai

cc @Borda @tchaton @justusschock @awaelchli @carmocca @ninginthecloud @daniellepintz @rohitgr7

[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!

[sjgosai]

Hello! I think this would be a really useful feature. It would improve the flexibility of setup in combination with the Trainer, as there is no other obvious solution for occasionally dynamic graph building. I haven't explored all of the hooks the trainer calls, but I'm certain there will be other use cases for skipping hooks in the trainer.

[vamaral1]

I had the same issue as reported in the PL forums but used a slightly different work-around which I'll provide here as an alternative solution. The issue revolved around PL trying to load the weights to the non-existing MyPytorchModule. Keep in mind it doesn't conform to the suggested use of callbacks

dm = MyDataModule()
model = MyNet()
trainer = pl.Trainer(callbacks=[ModuleInitializer()])
trainer.fit(model, dm)
trained_model = MyNet.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)

where

class ModuleInitializer(Callback):
    """Callback for dynamically initializing the model according to the learned data transformation
    from pl.LightningDataModule.setup. Called right after pl.LightningDataModule.setup"""

    def setup(
        self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: Optional[str] = None
    ):
        if stage == "fit":
            dim = trainer.datamodule.get_dim()
            pl_module.initialize(dim)
            pl_module.hparams["dim"] = dim

and

class MyNet(pl.LightningModule):

    def __init__(
        self,
        lr: float = 1e-3,
        dim = None
    ):
        super().__init__()
        self.lr = lr
        self.save_hyperparameters()
        if dim:
            self.initialize(dim)

    def initialize(self, dim):
        """Dynamically initialize the model according to the learned data transformation"""
        self.my_pytorch_module = MyPytorchModule(dim)
    ...

When the model is initialized for training it'll get the dimensionality from the ModuleInitializer callback and save those as hyperparameters. When the model is being loaded, PL passes in the saved hyperparemeters to the __init__ of the PL module which will allow the MyPytorchModule to be initialized and the saved weights loaded.

[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!

[ananthsub]

@awaelchli #6420 is a related view of this

[carmocca]

@awaelchli I would not pursue this as the described alternative is the natural solution for all Python users.

Doing this limits the flexibility in how they are skipped and requires knowing that .skip_hooks exists.

One more argument against is that "loop customization" is the natural vehicle for this. Albeit more flexible and more complex