Hook for Fully Formed Checkpoints

[dcharatan]

I would like to create a hook that automatically uploads checkpoints to the cloud (e.g., AWS, Azure) when they're created. I tried using on_save_checkpoint roughly like this:

def on_save_checkpoint(self, trainer: pl.Trainer, pl_module: pl.LightningModule, checkpoint: Dict[str, Any]) -> dict:
    checkpoint_bytes = io.BytesIO()
    torch.save(checkpoint, checkpoint_bytes)
    # Upload the BytesIO somewhere...

However, states for optimizers, schedulers, AMP, etc. are added after on_save_checkpoint hooks are called. Is there an elegant way to create a hook that receives the fully formed checkpoint state?

[rohitgr7]

hey @dcharatan !

I'd rather suggest using the remote filesystems. You can also specify the remote path inside ModelCheckpoint.

or use CheckpointIO plugin.

[dcharatan]

Thanks, this seems like exactly what I was looking for!