checkpoint every module in a different ckpt file

[malfonsoarquimea]

Hi!
I am currently working on a project where I would like to checkpoint my model in separated pieces.
My model has a backbone composed by:

• a backbone, which is also composed by 3 modules
• several heads, each one being a module
  I would like to save one ckpt with the backbone and one ckpt per head. I understand that I should create a custom callback inheriting from ModelCheckpoint and then modifying on_save_checkpoint, I am not really aware of how to do it.
  on_save_checkpoint is defined as:

Another solution would be to modify my lightning module to load the ckpt when the training ends as a dict and then save each subpart as a ckpt using torch.save(), but I understand that this solution is much less elegant.

Any suggestions? Thanks in advance!

[rohitgr7]

I'd suggest using checkpoint_io plugin for your use-case.

[malfonsoarquimea]

First of all, thanks very much for your answer!
I have some doubts regarding the checkpoint_io plugin
If I understood correctly, I should modify the method save_checkpoint() on my custom version of CheckpointIO(). My doubt (probably a noob doubt) is how this will interact with the ModelCheckpoint callback I am using in my training loop. I am a bit confused about callbacks and hooks and how they interact with each other, so if you could provide some insight on this particular topic or redirect me to some source of knowledge to understand how the callbacks and hooks actually work I would be really happy!

PD: Do you have any example of a custom CheckpointIO?
Thanks in advance!

[rohitgr7]

the checkpoint-related hooks in model_checkpoint are usually used to manipulate the checkpoint itself, more like what to checkpoint and add/delete/modify stuff in it. The CheckpointIO is built for users to give them the ability to modify how the checkpoint is saved and how it's being loaded.

for your use-case:

from pytorch_lightning.plugins.io import CheckpointIO

class MyCheckpointIO(CheckpointIO):
    def save_checkpoint(self, checkpoint, path, ...):
        # break the checkpoint['state_dict'] and store it separately

    def load_checkpoint(path, ...):
        # load the checkpoint and reconstruct it


custom_checkpoint_io = MyCheckpointIO()
trainer = Trainer(... plugins=[custom_checkpoint_io])

you can check the default one here for reference: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/io/torch_plugin.py

we also released a blog post recently that you can check: https://devblog.pytorchlightning.ai/lightning-checkpointio-plugin-f8663d8bc5e8

[malfonsoarquimea]

Thank you very much! I found really useful both things, the blog and the example. I managed to successfully save and load my ckpts as I wanted. Thanks very much!
Is there a way to update the state of the discussion to solved?

[rohitgr7]

Is there a way to update the state of the discussion to solved?

you can click on Mark as answer above in the first comment.

Also if possible can you share that piece of code if its possible for others looking for it in future?

[malfonsoarquimea]

For sure, I only implemented a dummy version where my model is structured in the base model and last layer. The state dict of the checkpoint contains the following keys

With that in mind, I created a custom IO plugin to save 2 ckpts, one with all the layers in the base model (contains the word "base" in the layer name) and the ones in the last layer (does not contain the word "base" in the layer name)

and the output of my training script is: