[RFC] Create a ModelCheckpointBase callback

[ananthsub]

🚀 Feature

Create a ModelCheckpointBase callback, and have the existing checkpoint callback extend it

Motivation

The model checkpoint callback is growing in complexity. Features that have been recently added or will soon be proposed:

• Checkpoint every n train steps: [feat] Support iteration-based checkpointing in model checkpoint callback #6146
• Checkpointing using a time-based interval during training: Support time-based checkpointing trigger #6286
• Checkpointing on train end to fix this hack
  https://github.com/PyTorchLightning/pytorch-lightning/blob/680e83adab38c2d680b138bdc39d48fc35c0cb58/pytorch_lightning/trainer/training_loop.py#L152-L163

The decision was made in #6146 to keep these triggers mutually exclusive, at least based on the phase they run in. Why? It's very hard to get the state management right. For instance, the monitor might be added for something that's available only during validation, but the checkpoint callback is configured to run during training too, and crashes when it tries to look up the monitor key in the available metrics for tracking. Tracking top-K models and scores is another huge pain. Supporting multiple monitor metrics on top of this is another beast.

cc @Borda @carmocca @awaelchli @ninginthecloud @jjenniferdai @rohitgr7

Pitch

Move the existing logic for the following into a base class:

• Core saving functionality
• Top-K model management
• formatting checkpoint names
• validation (though sub-classes should override this)

And have thin wrappers on top which extend this class and implement callback hook(s) for when to save the checkpoint.

Alternatives

The checkpoint callback gets bigger and bigger as we add more features to it.

Additional context

[dalek-who]

Related to this issue, I want to implement a logic: saving the best checkpoint's valid predict result during saving best checkpoint, and I want to save it in another file rather than in checkpoint's binary file. Right now, I can't know who calls LightningModule.on_save_checkpoint (can be best, best-k, lask or during exception).
If I implement it in LightningModule.on_save_checkpoint, I need a flag showing who save it; but I think the better way is a ModelCheckpointBase with some hooks like on_save_checkpoint_{best, best-k, last}

[carmocca]

Previous discussion: #4335 (comment)

Should we close that one in favor of this one? @ananthsub

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[ananthsub]

I have a different take on this now: which is extending the framework is relying on the class hierarchy to determine the intent of the callback. This is used by the Trainer especially to determine what callbacks are enabled by default. This is used for:

• checkpointing
• progress bar
• model summary

I think for these, we ought to create an empty base class. This way, users don't have to worry about changes to the concrete implementations also offered by the framework. As an example:

class BaseModelCheckpoint(Callback):
    pass

class ModelCheckpoint(BaseModelCheckpoint):
    # existing code today


class MyCustomModelCheckpoint(BaseModelCheckpoint)

Essentially, someone should be able to extend BaseModelCheckpoint, not ModelCheckpoint, which is completely empty, and customize this however they see fit without worrying about keeping their code in sync with the framework's changes to ModelCheckpoint.

@jjenniferdai @carmocca

[carmocca]

What would be the advantage internally by offering this base empty class?

[jjenniferdai]

custom checkpoint callbacks (AnyCustomUserCheckpointCallback(BaseModelCheckpoint)) don't have to inherit ModelCheckpoint logic and can still get detected and rearranged to execute last in callback_connector

[carmocca]

I am in favor of this.

I think it's becoming increasingly important that ModelCheckpoint gets split up. It has gotten bloated with multiple arguments that have chaotic or unintended interactions and make edge cases very confusing. For example:

• best_model_path does not retrieve the path to the best monitor checkpoint file #12485
• Raise a WARNING when someone tried to load the best checkpoint when one has not been set. #12501
• Fix interaction with save_last and every_n_epochs #12391
• Just look at the interactions of every_n_epochs: https://github.com/PyTorchLightning/pytorch-lightning/blob/3bcaed52454f3e6c3bce5513032e34302e5b1bb6/pytorch_lightning/callbacks/model_checkpoint.py#L117-L131

We couldn't do this in the past because we didn't support multiple instances of the same callback. Related: #4335 where I propose a split. We would need to revisit it because that discussion is old