Adding a warmup period to EarlyStopping and ModelCheckpoint

[rgmyr]

🚀 Feature

Add an optional warmup period for EarlyStopping and ModelCheckpoint callbacks.

Motivation

Sometimes the metric you want to monitor can take a number of epochs to stabilize and become meaningful.

For example: with GANs, you might want to monitor and minimize G's loss, but usually it starts out unreasonably low because it's based on the output of D, which hasn't learned anything about discriminating yet.

Pitch

I'd like to have this result in the callbacks having no effect for the first 10 epochs:

early_stop = EarlyStopping(..., warmup=10)
model_checkpoint = ModelCheckpoint(..., warmup=10)

Alternatives

I added this option through inheritance:

from pytorch_lightning.callbacks import EarlyStopping

class EarlyStoppingWithWarmup(EarlyStopping):
    """
    EarlyStopping, except don't watch the first `warmup` epochs.
    """
    def __init__(self, warmup=10, **kwargs):
        super().__init__(**kwargs)
        self.warmup = warmup

    def on_validation_end(self, trainer, pl_module):
        pass

    def on_epoch_end(self, trainer, pl_module):
        if trainer.current_epoch < self.warmup:
            return
        else:
            super()._run_early_stopping_check(trainer, pl_module)

Additional context

Edit: couldn't get this working at first, but got it figured out after upgrading my pl version. I would be happy to PR something like this if anyone can provide guidance on where to add it.

[rohitgr7]

Your code is working for me. Also make it trainer.current_epoch < self.warmup: since current_epoch starts from 0.

[rgmyr]

@rohitgr7 Thanks, was just editing my post as you said that. I've got it working now.

[s-rog]

for early stopping you can set min_epochs in trainer
https://pytorch-lightning.readthedocs.io/en/stable/trainer.html#min-epochs

for checkpoints, if you set save_top_k in ModelCheckpoint then the initial "bad" checkpoints will be overwritten

[rohitgr7]

yeah min_epochs will do the trick here but with val_check_interval != 1.0 it might not. Let's say I have a very big dataset and want to check with val_check_interval=0.5 with a warmup=7, then min_epochs won't work here. I think warmup should be specific to number of times early_stopping is called.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[thesofakillers]

@rohitgr7 I know this is 3 years late, but you can use min_steps instead of min_epochs in that case probably