Repository: Lightning-AI/pytorch-lightning
Pull request: Fixes resuming checkpoints rerunning last epoch
Fixes #850 


Problem statement:
Epoch end checkpoint restarts previous epoch
## 🐛 Bug

If restarting the training and reloading the model, the epoch that the checkpoint had just completed is restarted rather than beginning the next.

### Expected behavior

When a checkpoint upon epoch end is saved, restarting it should resume its state and start the next epoch.

Hints from issue discussion:
https://github.com/PyTorchLightning/pytorch-lightning/blob/edd4a87fb0e976e004a82c7a3491598ab65cafbd/pytorch_lightning/trainer/training_io.py#L389

This seems as simple as replacing the line above with ```self.current_epoch = checkpoint['epoch'] + 1``` since the checkpointers save at the end of validation and the main loop runs from the current epoch.

We should probably also increase the global step by 1 since this happens after saving the checkpoint.
https://github.com/PyTorchLightning/pytorch-lightning/blob/edd4a87fb0e976e004a82c7a3491598ab65cafbd/pytorch_lightning/trainer/training_io.py#L388

https://github.com/PyTorchLightning/pytorch-lightning/blob/edd4a87fb0e976e004a82c7a3491598ab65cafbd/pytorch_lightning/trainer/training_loop.py#L411

https://github.com/PyTorchLightning/pytorch-lightning/blob/edd4a87fb0e976e004a82c7a3491598ab65cafbd/pytorch_lightning/trainer/training_loop.py#L430 

I'll add a test in whilst I do it. 

@williamFalcon @Borda Any other thoughts if I put a PR in with this? The test should presumably go in ``test_trainer.py``?

That works when saving at epoch end but there's many cases of saving during an epoch as well (e.g. for a very large dataset). Both the epoch number and the global step are technically correct upon saving mid or end epoch but when resuming, the loop starts at the beginning. The best solution is to make the resume reliably restart precisely where in the loop it left off. I'm not that familiar with this code but guess one should then also save the batch_idx.
I can think of at least one way to do it, although it's not ideal:
- Use the global step as stored and when resuming, load this into a (hidden?) trainer variable which is checked every training batch. 
- Skip each batch until the batch number is above this saved variable, at which point we can set it to ``None`` or something. 
- Run batches as normal

Two Problems I see:
- When the data set is shuffled there would be no guarantee of seeing only new samples after the reload without somehow 'resuming' the dataloaders. 
- Any calls that are usually made during the batch wouldn't happen. For example, the tqdm update calls would need to be faked so it didn't appear to end the epoch early in the progress bars.

Commit message: Formatting

Patch to review:
diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py
index 2ee77ada4c5c4..7ae2d53609758 100644
--- a/pytorch_lightning/trainer/training_io.py
+++ b/pytorch_lightning/trainer/training_io.py
@@ -125,6 +125,7 @@ def __init__(self):
         self.early_stop_callback = None
         self.lr_schedulers = None
         self.optimizers = None
+        self.num_training_batches = None
 
     def get_model(self):
         is_dp_module = isinstance(self.model, (LightningDistributedDataParallel,
@@ -305,10 +306,9 @@ def restore(self, checkpoint_path, on_gpu):
         self.restore_training_state(checkpoint)
 
     def dump_checkpoint(self):
-
         checkpoint = {
-            'epoch': self.current_epoch,
-            'global_step': self.global_step
+            'epoch': self.current_epoch + 1,
+            'global_step': self.global_step + 1,
         }
 
         if self.checkpoint_callback is not None and self.checkpoint_callback is not False:
@@ -388,6 +388,14 @@ def restore_training_state(self, checkpoint):
         self.global_step = checkpoint['global_step']
         self.current_epoch = checkpoint['epoch']
 
+        # Deals with peculiarity of different global step for odd vs even num_training_batches
+        if abs((self.global_step + 1) % self.num_training_batches) > 1:
+            warnings.warn(
+                "You're resuming from a checkpoint that ended mid-epoch. "
+                "This can cause unreliable results if further training is done, "
+                "consider using an end of epoch checkpoint. "
+            )
+
         # restore the optimizers
         optimizer_states = checkpoint['optimizer_states']
         for optimizer, opt_state in zip(self.optimizers, optimizer_states):
diff --git a/tests/test_restore_models.py b/tests/test_restore_models.py
index 92ab71364ccae..2bd033b09756f 100644
--- a/tests/test_restore_models.py
+++ b/tests/test_restore_models.py
@@ -214,8 +214,8 @@ def test_dp_resume(tmpdir):
     trainer.is_slurm_managing_tasks = True
     result = trainer.fit(model)
 
-    # track epoch before saving
-    real_global_epoch = trainer.current_epoch
+    # track epoch before saving. Increment since we finished the current epoch, don't want to rerun
+    real_global_epoch = trainer.current_epoch + 1
 
     # correct result and ok accuracy
     assert result == 1, 'amp + dp model failed to complete'
@@ -282,7 +282,8 @@ def test_cpu_restore_training(tmpdir):
     # fit model
     trainer = Trainer(**trainer_options)
     result = trainer.fit(model)
-    real_global_epoch = trainer.current_epoch
+    # Increment since we've finished the current epoch, don't want to rerun
+    real_global_epoch = trainer.current_epoch + 1
 
     # traning complete
     assert result == 1, 'amp + ddp model failed to complete'
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
index 06cd8a7de93cd..c4b5098e905ad 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -410,6 +410,74 @@ class CurrentTestModel(
         tutils.run_prediction(dataloader, trainer.model)
 
 
+def test_resume_from_checkpoint_epoch_restored(tmpdir):
+    """Verify resuming from checkpoint runs the right number of epochs"""
+    import types
+
+    tutils.reset_seed()
+
+    hparams = tutils.get_hparams()
+
+    def new_model():
+        # Create a model that tracks epochs and batches seen
+        model = LightningTestModel(hparams)
+        model.num_epochs_seen = 0
+        model.num_batches_seen = 0
+
+        def increment_epoch(self):
+            self.num_epochs_seen += 1
+
+        def increment_batch(self, _):
+            self.num_batches_seen += 1
+
+        # Bind the increment_epoch function on_epoch_end so that the
+        # model keeps track of the number of epochs it has seen.
+        model.on_epoch_end = types.MethodType(increment_epoch, model)
+        model.on_batch_start = types.MethodType(increment_batch, model)
+        return model
+
+    model = new_model()
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_epochs=2,
+        train_percent_check=0.65,
+        val_percent_check=1,
+        checkpoint_callback=ModelCheckpoint(tmpdir, save_top_k=-1),
+        logger=False,
+        default_save_path=tmpdir,
+        early_stop_callback=False,
+        val_check_interval=0.5,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.fit(model)
+
+    training_batches = trainer.num_training_batches
+
+    assert model.num_epochs_seen == 2
+    assert model.num_batches_seen == training_batches * 2
+
+    # Other checkpoints can be uncommented if/when resuming mid-epoch is supported
+    checkpoints = [
+        # os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_0.ckpt"),
+        os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_0_v0.ckpt"),
+        # os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_1.ckpt"),
+        os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_1_v0.ckpt"),
+    ]
+
+    for check in checkpoints:
+        next_model = new_model()
+        state = torch.load(check)
+
+        # Resume training
+        trainer_options['max_epochs'] = 4
+        new_trainer = Trainer(**trainer_options, resume_from_checkpoint=check)
+        new_trainer.fit(next_model)
+        assert state['global_step'] + next_model.num_batches_seen == training_batches * 4
+
+
 def test_multiple_test_dataloader(tmpdir):
     """Verify multiple test_dataloader."""
     tutils.reset_seed()


Code review comments:
diff_hunk: @@ -388,6 +388,14 @@ def restore_training_state(self, checkpoint):
         self.global_step = checkpoint['global_step']
         self.current_epoch = checkpoint['epoch']
 
+        # Deals with peculiarity of different global step for odd vs even num_training_batches
+        if abs((self.global_step + 1) % self.num_training_batches) > 1:
path: pytorch_lightning/trainer/training_io.py
text: @author:
Is this known?

@user1:
the `abs` looks suspicious...

@author:
The check should be `!=0` if the global step matched for odd vs even number of training steps, so if that's worked out then we shouldn't need it. Unless it's intended, in which case I think just ``(self.global_step + 1) % self.num_training_batches`` is sufficient since global step matches num_training_batches in the even case. 

Thinking about it, I should probably just remove the abs, it is fine without. 

@author:
I've removed the `abs` and made sure it handled accumulated batches properly. Can you think of anything else that might change the global step? 

Currently the only test that throws this warning now is ``test_restore_models/test_dp_resume`` since this changes the percentage of train data used when resuming in a new trainer. Not much we can about that.