How to save last k checkpoints?

[hiyyg]

In pl's checkpointer, how can we save the last k checkpoints with a pre-defined checkpointing interval (such as every 5000 iterations or every 5 epochs)? It seems the save_last option only saves the last 1 checkpoint.

[tchaton]

Hey @hiyyg,

You can use ModelCheckpoint(save_top_k=X).

[rohitgr7]

maybe you can set ModelCheckpoint(save_top_k=k, mode='max', monitor='time_log') and log the timestamp. This will consider timestamp as the metric to monitor for keeping the best k checkpoints.

[hal-314]

This workaround doesn't work in PL 1.5.4. There is "step" key as available key to monitor but "_validate_monitor_key" function (Line 377 in model_checkpoint.py) prevents it's use. Setting every_n_train_steps and every_n_epochs doesn't work as save_top_k needs a monitor different to None

My workaround is to use a custom model checkpoint class and then call it as ModelCheckpointWorkaround(save_top_k=k, mode='max', monitor='step') where

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint as PLModelCheckpoint


class ModelCheckpointWorkaround(PLModelCheckpoint):
    """Like pytorch_lightning.callbacks.ModelCheckpoint but allowing saving last top k checkpoints.
    See https://github.com/PyTorchLightning/pytorch-lightning/discussions/10669
    """
    def _validate_monitor_key(self, trainer: pl.Trainer) -> None:
        if self.monitor in ("step", "epoch"):  # Lightning check in 1.5.x doesn't contemplate them.
            return
        super()._validate_monitor_key(trainer)

I think it would be nice to have this functionality inside Lightning. In my company, we use steps instead of epochs as we have a lot of training data. So, we want to validate and save checkpoints several times in one training epoch. For us, it's not possible to use a neither train loss (very noise) or validation loss / metrics (we want to save several checkpoints between every validation).

[hal-314]

@rohitgr7 Do you think that I should open an issue so Lightning is able to support trainings based on steps instead of epochs?

[rohitgr7]

did you log global_step?

def training_step(self, ...):
    ...
    self.log('monitoring_step', self.global_step)

ModelCheckpoint(..., monitor='monitoring_step')

[hal-314]

I wasn't logging manually the global step. I'm using W&B logger and it appears automatically as trainer/global_step. However, ModelCheckpoint isn't aware of it.

Logging manually as you suggested works. Thank you

[rohitgr7]

glad it worked!
yeah if you are monitoring something, it needs to be logged. Wandb might be a special case here. we need the codebase to be aligned for all kinds of loggers.