Duplicate epochs when calling .fit() twice

[carmocca]

🐛 Bug

To Reproduce

def test_bug(tmpdir):
    epochs = []
    
    class TestModel(BoringModel):
        def on_epoch_end(self):
            epochs.append(self.current_epoch)
    
    trainer = Trainer(
        max_epochs=2,
        limit_train_batches=1,
        limit_val_batches=1,
        default_root_dir=tmpdir,
        checkpoint_callback=False,
        logger=False,
        weights_summary=None,
        progress_bar_refresh_rate=0,
    )
    trainer.fit(TestModel())
    trainer.max_epochs=4
    trainer.fit(TestModel())

    assert epochs == list(range(4))
    # AssertionError [0, 1, 1, 2, 3] != [0, 1, 2, 3]

Expected behavior

Assertion does not fail

Environment

Current master

cc @tchaton @Borda

[carmocca]

The epoch number is generated here:

https://github.com/PyTorchLightning/pytorch-lightning/blob/239347435029c0a02b305201ebbfa39d62746ca8/pytorch_lightning/trainer/trainer.py#L511

which assumes that self.current_epoch has not run yet. This assumption is not correct when fit is run twice because the epoch number is not increased after training ends.

When using Trainer(resume_from_checkpoint=...) this issue does not appear due to this piece:

https://github.com/PyTorchLightning/pytorch-lightning/blob/239347435029c0a02b305201ebbfa39d62746ca8/pytorch_lightning/trainer/connectors/checkpoint_connector.py#L273-L279

So the solution would be to increase the epoch number at the end of training_loop.on_train_end()? (This would break backwards compatibility)

[ananthsub]

@carmocca or should we reset these parameters in the trainer teardown?

[carmocca]

I'd say it's more natural to do it on_train_end because it has nothing to do with test and teardown is for fit or test

[pierresegonne]

Any opinion on what would be the recommended way of resetting the current_epoch to call fit twice?

Is something along those lines

trainer.fit(model, datamodule)
model.trainer.current_epoch = 0
trainer.fit(model, datamodule)

safe ?

[carmocca]

Yes, it should be safe.

[edenlightning]

@carmocca what is left TODO here?

[carmocca]

Everything, the bug is not fixed 🙂

There is a reproduction test at the top. We just need to make our minds about the best solution. Context here: #5007 (comment)

[carmocca]

Status update: WIP - tackling other related issues first. Need this for fault-tolerance

[carmocca]

Status update: Blocked by merging #8477 and enabling restoring the ckpt progress tracking state by default.

[fschiffers]

Any opinion on what would be the recommended way of resetting the current_epoch to call fit twice?

Is something along those lines

trainer.fit(model, datamodule)
model.trainer.current_epoch = 0
trainer.fit(model, datamodule)

safe ?

I think the current_epoch can no longer be set in the trainer but must be set in the fit_loop itself.

[carmocca]

Correct. You'll need to do trainer.fit_loop.current_epoch = 0