How to checkpoint submodules

[dtch1997]

By default pytorch lightning Trainer saves checkpoint for the full model. However I would like to checkpoint manually for individual submodules.

Eg. I might have an image classifier that predicts labels and a image caption generator that generates captions from images; both of them use the same CNN encoder. I might want to do something like pre-train with the classification task (saving checkpoints) and then load only the trained encoder (minus the prediction head) for use in the downstream caption task. Furthermore I might have several backbones that I want to test out for the encoder/

What's the idiomatic way to do this in PyTorch Lightning?

[justusschock]

Hi,

to do this, you probably have to provide your own ModelCheckpoint Callback

[dtch1997]

thanks @justusschock . Are there code snippets of custom checkpoint callbacks that I can refer to?

[justusschock]

None that I know of, to be honest.

[dtch1997]

Alright thanks. I think I'll stick to using torch.save and torch.load for now then.