EarlyStopping not updating it's value after resuming training

[MaugrimEP]

Bug description

Early stopping callback does not update its best value after resuming training from a checkpoint.

What version are you seeing the problem on?

v2.0

How to reproduce the bug

ckpt_path = "some_path/66_best_model.ckpt"
es = EarlyStopping(
		monitor  ="valid/step/loss",
		min_delta=0,
		patience =100,
		verbose  =True,
		mode     ="min",
        check_on_train_epoch_end=False,
	)

trainer = pl.Trainer(callbacks=[es])
trainer.fit(model, datamodule=datamodule, ckpt_path= ckpt_path)

Error messages and logs

Hi,
I have an EarlyStopping and ModelCheckpoint on the same metrics in my callbacks. Also, I continue an existing training.

It seems like EarlyStopping does not update its best metrics if we load a checkpoint (retraining).

I got this output:

Epoch 149, global step 50700: 'epoch' reached 149.00000 (best 149.00000), saving model to '/dlocal/run/19197/_model_save/149_duration_model.ckpt' as top 100
Epoch 150, global step 51038: 'valid/step/loss' reached 0.00266 (best 0.00266), saving model to '/dlocal/run/19197/_model_save/150_best_model.ckpt' as top 1
Epoch 151, global step 51376: 'valid/step/loss' was not in top 1
Epoch 152, global step 51714: 'valid/step/loss' was not in top 1
Epoch 153, global step 52052: 'valid/step/loss' was not in top 1
Epoch 154, global step 52390: 'valid/step/loss' was not in top 1
Epoch 155, global step 52728: 'valid/step/loss' was not in top 1
Epoch 156, global step 53066: 'valid/step/loss' was not in top 1
Epoch 157, global step 53404: 'valid/step/loss' was not in top 1
Epoch 158, global step 53742: 'valid/step/loss' was not in top 1
Epoch 159, global step 54080: 'valid/step/loss' was not in top 1
Epoch 159, global step 54080: 'epoch' reached 159.00000 (best 159.00000), saving model to '/dlocal/run/19197/_model_save/159_duration_model.ckpt' as top 100
Generation Loop: 100%|██████████| 100/100 [05:37<00:00,  3.38s/it]it]2:26<03:13,  3.40s/it]
Generation Loop: 100%|██████████| 100/100 [05:38<00:00,  3.38s/it]
Generation Loop: 100%|██████████| 100/100 [05:38<00:00,  3.39s/it]
Generation Loop: 100%|██████████| 100/100 [05:39<00:00,  3.40s/it]
Epoch 160, global step 54418: 'valid/step/loss' was not in top 1
Epoch 161, global step 54756: 'valid/step/loss' was not in top 1
Epoch 162, global step 55094: 'valid/step/loss' was not in top 1
Epoch 163, global step 55432: 'valid/step/loss' was not in top 1
Epoch 164, global step 55770: 'valid/step/loss' was not in top 1
Epoch 165, global step 56108: 'valid/step/loss' was not in top 1
[rank: 0] Monitored metric valid/step/loss did not improve in the last 100 records. Best score: 0.003. Signaling Trainer to stop.
[rank: 3] Monitored metric valid/step/loss did not improve in the last 100 records. Best score: 0.003. Signaling Trainer to stop.
[rank: 1] Monitored metric valid/step/loss did not improve in the last 100 records. Best score: 0.003. Signaling Trainer to stop.
[rank: 2] Monitored metric valid/step/loss did not improve in the last 100 records. Best score: 0.003. Signaling Trainer to stop.
Epoch 166, global step 56446: 'valid/step/loss' was not in top 1

As we can see, in epoch 150 the model indeed improves.
But in epoch 165 the model stops (Early stopping with patience of 100), which should not be the case.
In fact, I resumed training from 66_best_model.ckpt. This ckpt file is in a specific path outside my project for reasons.
It looks like the best metrics did not change from this ckpt.

Environment

No response

More info

No response

cc @carmocca @awaelchli

[ryan597]

Think I have pinpointed the issue for this.
This happens when you have run trainer.fit, an early stop is triggered, and you try to run trainer.fit with this checkpoint path in the same program run. If you run in a new python instance this issue doesn't occur, or if you re-initialize the Trainer.

The problem is that when the early stop is triggered, the trainer.should_stop flag is not reset, and upon loading a new checkpoint the flag is not reset either. It is only reset when the trainer is re-initialized.

As you would expect it to be reset (if you call fit, you want to continue training anyways and you have probably changed something in between), it makes sense to reset this flag in the trainer.fit.

 def fit(
        ...
        self.state.fn = TrainerFn.FITTING
        self.state.status = TrainerStatus.RUNNING
        self.training = True
        self.should_stop = False
        ...
)

Full Reproducible

using the autoencoder example file

from os import path
from typing import Optional, Tuple

import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, random_split

from lightning.pytorch import callbacks, cli_lightning_logo, LightningDataModule, LightningModule, Trainer
from lightning.pytorch.cli import LightningCLI
from lightning.pytorch.demos.mnist_datamodule import MNIST
from lightning.pytorch.utilities import rank_zero_only
from lightning.pytorch.utilities.imports import _TORCHVISION_AVAILABLE

if _TORCHVISION_AVAILABLE:
    import torchvision
    from torchvision import transforms
    from torchvision.utils import save_image

DATASETS_PATH = path.join(path.dirname(__file__), "..", "..", "Datasets")

class ImageSampler(callbacks.Callback):
    def __init__(
        self,
        num_samples: int = 3,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        norm_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: int = 0,
    ) -> None:
        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
            raise ModuleNotFoundError("You want to use `torchvision` which is not installed yet.")

        super().__init__()
        self.num_samples = num_samples
        self.nrow = nrow
        self.padding = padding
        self.normalize = normalize
        self.norm_range = norm_range
        self.scale_each = scale_each
        self.pad_value = pad_value

    def _to_grid(self, images):
        return torchvision.utils.make_grid(
            tensor=images,
            nrow=self.nrow,
            padding=self.padding,
            normalize=self.normalize,
            value_range=self.norm_range,
            scale_each=self.scale_each,
            pad_value=self.pad_value,
        )

    @rank_zero_only
    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        if not _TORCHVISION_AVAILABLE:
            return

        images, _ = next(iter(DataLoader(trainer.datamodule.mnist_val, batch_size=self.num_samples)))
        images_flattened = images.view(images.size(0), -1)

        # generate images
        with torch.no_grad():
            pl_module.eval()
            images_generated = pl_module(images_flattened.to(pl_module.device))
            pl_module.train()

        if trainer.current_epoch == 0:
            save_image(self._to_grid(images), f"grid_ori_{trainer.current_epoch}.png")
        save_image(self._to_grid(images_generated.reshape(images.shape)), f"grid_generated_{trainer.current_epoch}.png")


class LitAutoEncoder(LightningModule):
    def __init__(self, hidden_dim: int = 64, learning_rate=10e-3):
        super().__init__()
        self.save_hyperparameters()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 3))
        self.decoder = nn.Sequential(nn.Linear(3, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 28 * 28))

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z)

    def training_step(self, batch, batch_idx):
        loss = self._common_step(batch, batch_idx, "train")
        self.log("loss", loss, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        self._common_step(batch, batch_idx, "val")

    def test_step(self, batch, batch_idx):
        self._common_step(batch, batch_idx, "test")

    def predict_step(self, batch, batch_idx, dataloader_idx=None):
        x = self._prepare_batch(batch)
        return self(x)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

    def _prepare_batch(self, batch):
        x, _ = batch
        return x.view(x.size(0), -1)

    def _common_step(self, batch, batch_idx, stage: str):
        x = self._prepare_batch(batch)
        loss = F.mse_loss(x, self(x))
        self.log(f"{stage}_loss", loss, on_step=True)
        return loss


class MyDataModule(LightningDataModule):
    def __init__(self, batch_size: int = 32):
        super().__init__()
        dataset = MNIST(DATASETS_PATH, train=True, download=True, transform=transforms.ToTensor())
        self.mnist_test = MNIST(DATASETS_PATH, train=False, download=True, transform=transforms.ToTensor())
        self.mnist_train, self.mnist_val = random_split(dataset, [55000, 5000])
        self.batch_size = batch_size

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size, num_workers=10)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=self.batch_size, num_workers=10)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=10)

    def predict_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=10)


def cli_main():
    checkpoint_cb = callbacks.ModelCheckpoint(dirpath="checkpoints", filename="test", monitor="loss", verbose=True)
    stopping_cb = callbacks.EarlyStopping(monitor="loss", verbose=True)
    cli = LightningCLI(
        LitAutoEncoder,
        MyDataModule,
        seed_everything_default=1234,
        run=False,  # used to de-activate automatic fitting.
        trainer_defaults={"callbacks": [ImageSampler(), checkpoint_cb, stopping_cb], "max_epochs": 100},
        save_config_kwargs={"overwrite": True},
    )
    cli.trainer.fit(cli.model, datamodule=cli.datamodule)

    print("\n\n\n CONTINUING FROM CKPT\n\n\n")
    # cli.trainer.should_stop = False  #### UNCOMMENT THIS LINE TO FIX BUG
    cli.trainer.fit(cli.model, ckpt_path="checkpoints/test.ckpt", datamodule=cli.datamodule)

if __name__ == "__main__":
    cli_main()

[awaelchli]

@ryan597 Would it resolve the issue if we set self.should_stop = False at the beginning of the fit method in Trainer?

[ryan597]

Yes it does @awaelchli (see the code snippet with the suggested fix).
I'll make the PR shortly for this if that's all good :)