Training: save checkpoints by iterations instead of by time

[trivoldus28]

Sometimes it is useful to save/backup checkpoints by iterations instead of by some time durations (e.g., every 15 minutes) to be able to quickly check the model performance across iterations. DeepEM does this.

Technically it is possible to query the checkpoint files in the bucket but it's just a bit too complex to handle in a simple script or cue spec.

[trivoldus28]

I've been testing this

@typeguard.typechecked
def get_checkpointing_callbacks(
    log_dir: str,
    backup_every_n_secs: int | None = None,
    update_every_n_secs: int | None = None,
    backup_every_n_steps: int | None = None,
    update_every_n_steps: int | None = None,
) -> List[pl.callbacks.Callback]:  # pragma: no cover
    result = []  # type: List[pl.callbacks.Callback]

    if backup_every_n_secs is not None:
        callback = pl.callbacks.ModelCheckpoint(
            dirpath=log_dir,
            train_time_interval=datetime.timedelta(seconds=backup_every_n_secs),
            save_top_k=-1,
            save_last=False,
            monitor=None,
            filename="{epoch}-{step}-backup",
        )
        result.append(callback)

    if update_every_n_secs is not None:
        save_top_k = 1 if backup_every_n_secs is None else 0  # keep copy if backup is not enabled
        callback = pl.callbacks.ModelCheckpoint(
            dirpath=log_dir,
            train_time_interval=datetime.timedelta(seconds=update_every_n_secs),
            save_top_k=save_top_k,
            save_last=True,
            monitor=None,
            filename="{epoch}-{step}-current",
        )
        result.append(callback)

    if backup_every_n_steps is not None:
        callback = pl.callbacks.ModelCheckpoint(
            dirpath=log_dir,
            every_n_train_steps=backup_every_n_steps,
            save_top_k=-1,
            save_last=False,
            monitor=None,
            filename="{step}",
        )
        result.append(callback)

    if update_every_n_steps is not None:
        save_top_k = 1 if backup_every_n_steps is None else 0  # keep copy if backup is not enabled
        callback = pl.callbacks.ModelCheckpoint(
            dirpath=log_dir,
            every_n_train_steps=update_every_n_steps,
            save_top_k=save_top_k,
            save_last=True,
            monitor=None,
            filename="{step}",
        )
        result.append(callback)
    
    return result