How to ensure objects saved as model attributes are saved in the checkpoint file?

[KirillShmilovich]

Say I have a lightning model model = MyLightningModel() that contains an object that gets created and updated throughout training model.my_object. Upon loading the model from the checkpoint MyLightningModel.load_from_checkpoint(ckpt_path) I'm noticing the model.my_object attribute is not being saved and restored upon loading from checkpoint.

Is it possible to somehow ensure that these model attributes get saved in the checkpoint file and properly restored when loading the model from checkpoint?

Thanks in advance for your help!

[tchaton]

Dear @KirillShmilovich.

You could use the LightningModule on_save_checkpoint and on_load_checkpoint hooks.

class MyLightningModel(LightningModule):

    def on_save_checkpoint(self):
        return {"my_object": self.my_object}

    def on_load_checkpoint(self, state_dict):
        self.my_object = state_dict["my_object"]

However, pickling objets isn't always the best approach. A slightly better approach is

class MyLightningModel(LightningModule):

    def on_save_checkpoint(self):
        return {"my_object_state_dict": self.my_object.state_dict()}

    def on_load_checkpoint(self, state_dict):
        self.my_object = my_object_cls.from_state_dict(state_dict["my_object_state_dict"])