Add required states for resumed ModelCheckpoint GC

[ORippler]

What does this PR do?

Fixes #4911
Related: #5090

Currently, when resuming training the internal states required for continued ModelCheckpointing are neither saved nor restored. This leads to the fact that k new checkpoints are always generated due to this check. These new checkpoints are properly gced/compared to against, but the old ones are not.

Note that this PR does not handle overrides of monitor, dirpath or mode, as also referred to in #4911

Does your PR introduce any breaking changes? If yes, please list them.

It might be that resuming training fails now if it did not fail before, if the paths were changed in the mean time (refer also #4911). I did not check/test for this, but confirmed that resumption of GC now works properly

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃
cc @carmocca @awaelchli @ninginthecloud @jjenniferdai

[justusschock]

@ORippler thanks for this fix. To avoid regression again, we need a test for this.

Do you think the following does reflect this issue sufficiently (If so, feel free to take it and commit it directly to your branch):

def test_model_checkpoint_attributes(tmpdir):
    seed_everything()
    model = LogInTwoMethods()

    epochs = 2
    checkpoint_callback = ModelCheckpoint(monitor=None, dirpath=tmpdir, save_top_k=-1, save_last=True)
    trainer = Trainer(
        default_root_dir=tmpdir,
        callbacks=[checkpoint_callback],
        limit_train_batches=10,
        limit_val_batches=10,
        max_epochs=epochs,
        logger=False,
    )

    trainer.fit(model)

    checkpoint = torch.load(os.path.join(tmpdir, 'last.ckpt'))['callbacks'][checkpoint_callback.state_key]
    for k in ("best_models, kth_best_model_path", "kth_value", "last_model_path"):
        assert checkpoint[k] == getattr(checkpoint_callback, k)

[ORippler]

@ORippler thanks for this fix. To avoid regression again, we need a test for this.

Do you think the following does reflect this issue sufficiently (If so, feel free to take it and commit it directly to your branch):

def test_model_checkpoint_attributes(tmpdir):
    seed_everything()
    model = LogInTwoMethods()

    epochs = 2
    checkpoint_callback = ModelCheckpoint(monitor=None, dirpath=tmpdir, save_top_k=-1, save_last=True)
    trainer = Trainer(
        default_root_dir=tmpdir,
        callbacks=[checkpoint_callback],
        limit_train_batches=10,
        limit_val_batches=10,
        max_epochs=epochs,
        logger=False,
    )

    trainer.fit(model)

    checkpoint = torch.load(os.path.join(tmpdir, 'last.ckpt'))['callbacks'][checkpoint_callback.state_key]
    for k in ("best_models, kth_best_model_path", "kth_value", "last_model_path"):
        assert checkpoint[k] == getattr(checkpoint_callback, k)

How would this integrate with the different functionality tests for ModelCheckpoint the overall testing framework ? I see many tests checking whether attributes are written to the ckpt properly, but not whether they are loaded. For example here, we check for current_score being written to disk/to the ckpt, but the current on_load_checkpoint never sets this attribute again.

Is this not something we want to test also? Am a bit confused here.

[justusschock]

@ORippler that's true. I suppose for end-to-end testing you would have to extend that test by resuming it at a freshly created trainer and examining the properties there.

@carmocca do we want to test for different parametrizations of the callback here?

[ORippler]

@ORippler that's true. I suppose for end-to-end testing you would have to extend that test by resuming it at a freshly created trainer and examining the properties there.

@carmocca do we want to test for different parametrizations of the callback here?

I added your test and expanded it to check whether a freshly instantiated ModelCheckpoint also loads the properties.
This is still not a fully functional test though.

Off-Note: Do we have a test that compares for equivalence of results generated by one continuous training run and an interrupted one that is resumed by passing the checkpoint to trainer.fit ? This would be nice to have imo

[justusschock]

@ORippler we have https://github.com/PyTorchLightning/pytorch-lightning/blob/master/tests/trainer/test_trainer.py#L399 which doesn't check the results. However, this is on purpose as the checkpoint does not include any random state and thus continuing from the checkpoint doesn't have to yield the exact same results (different random states when using the global rng for example) until now (this is currently in development).

cc @tchaton to add a similar test once fault tolerance is ready

[tchaton]

Hey @ORippler,

Yes, we have multiple tests checking the weights are the same before and after for Fault Tolerance. Here they are: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/tests/utilities/test_auto_restart.py

[codecov]

Codecov Report

Merging #10995 (3d7994a) into master (e19d93f) will decrease coverage by 4%.
The diff coverage is 100%.

@@           Coverage Diff            @@
##           master   #10995    +/-   ##
========================================
- Coverage      92%      88%    -4%     
========================================
  Files         177      177            
  Lines       16502    16560    +58     
========================================
- Hits        15173    14604   -569     
- Misses       1329     1956   +627     

[awaelchli]

LGTM

[carmocca]

I just noticed an issue with doing this.

Since we save each "ModelCheckpoint" mode sequentally, these attributes will not be correct depending on the order if more than 1 mode triggers a save for the same global step:

https://github.com/PyTorchLightning/pytorch-lightning/blob/fe940e195dceb18eb9f3bd512cea56ae3405d464/pytorch_lightning/callbacks/model_checkpoint.py#L366-L373

Currently, a "top-k" checkpoint will not include the last_model_path path even if it's saved right after for this global step.

I'm not sure what would be the best solution here. I think we should start recommending multiple ModelCheckpoint instances as a best practice because these interactions between flags can be unintuitive.

cc @awaelchli @ananthsub @jjenniferdai
Related to #4335 and #11805 (comment)