[feat] Support iteration-based checkpointing in model checkpoint callback

[ananthsub]

What does this PR do?

Fixes #2534

This supports checkpointing periodically after a user-specified number of training batches to complement the end of validation epoch-based checkpointing currently supported. This is useful for large training runs and for purely iteration-based training.

We checkpoint after training batches because the model state doesn't change during validation. This means the training metrics will be used for the monitor keys if configured to checkpoint this way.

To do so:

• We add a new argument every_n_train_steps to the callback constructor which dictates the frequency with which we'll checkpoint. If every_n_train_steps = 0 then this feature is disabled.
• The period name is ambiguous with this addition: therefore we introduce a new argument every_n_val_epochs to be more clear as to when it's applied and mark period as deprecated
• We retain the existing functionality of checkpointing after each validation epoch
• These triggers for when to save the checkpoint must be mutually exclusive to avoid additional complexity with handling monitor values and top-K results across different phases (train, val)

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 update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
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 🙃

[pep8speaks]

Hello @ananthsub! Thanks for updating this PR.

• In the file pytorch_lightning/callbacks/model_checkpoint.py:

Line 235:13: W503 line break before binary operator
Line 290:13: W503 line break before binary operator
Line 291:13: W503 line break before binary operator
Line 292:13: W503 line break before binary operator

Comment last updated at 2021-03-11 16:58:30 UTC

[codecov]

Codecov Report

Merging #6146 (4065f42) into master (d1db604) will decrease coverage by 2%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #6146    +/-   ##
=======================================
- Coverage      94%     92%    -2%     
=======================================
  Files         161     161            
  Lines       11500   11512    +12     
=======================================
- Hits        10756   10547   -209     
- Misses        744     965   +221     

[carmocca]

Got interrupted during my review. Will re-review in another moment

[carmocca]

Should None be the default?

[ananthsub]

for every_n_batches? would we want every_n_epochs to be optional too?

one choice is to make both None by default, and then in the constructor, if neither are set, then we default to every_n_epochs=1 so it runs after each validation epoch.

And if someone wants to checkpoint after batches only, they can do ModelCheckpoint(..., every_n_batches=N) without needing to toggle off the every_n_epochs flag. What do you think?

[carmocca]

Yes, both optional. The question is whether to assume they are exclusive and we go with

make both None by default, and then in the constructor, if neither are set, then we default to every_n_epochs=1 so it runs after each validation epoch.

Or not, have every_n_epochs=1 by default and have the user toggle off if necessary.

Does allowing having both active make sense to you?

[ananthsub]

I think it's unlikely that both will be used together, but they aren't mutually exclusive. if you want to checkpoint every N batches, you can. if you want to checkpoint every M epochs, you can. if you set both, we'll do both

Since this regular int is more restrictive than Optional[int] wdyt about going with this approach for now, and if we get feedback that this is confusing, then we can make them both optional to maintain backwards compatibility?

[carmocca]

Sounds good to me. To recap:

• <0: Misconfiguration
• 0: disabled / acts as identity
• n: every n

both can be set at the same time

defaults:

• every_n_epochs=1
• every_n_steps=0

Correct?

[ananthsub]

correct!

[ananthsub]

@carmocca @awaelchli coming back to this, after looking at the complexity in supporting monitor and save_top_k, it makes more sense to keep them as mutually exclusive.

• So we can make both every_n_train_steps and every_n_val_epochs arguments as Optional[int]
• users toggle on whichever they're going to use
• we raise an exception if both are set and positive
• if neither are set, we default to every_n_train_steps = 0, every_n_val_epochs = 1
• <0: Misconfiguration
• 0: Disabled
• n: every n

[carmocca]

That's fine with me, especially if we support many callbacks

[ananthsub]

@carmocca @tchaton @SkafteNicki i think this is mostly ready for review now. the main question i have is whether we should be using the total batch idx or the global step in the check for when to save during training. depending on that, we can set the appropriate name for the parameter (every_n_batches vs every_n_steps)