Saving only student model in ModelCheckpoint

[ari9dam]

Hi

if my Pytorch Lightning model contains two modules:

self.teacher = AutoModelForSequenceClassification.from_pretrained(hparams["teacher_model_name_or_path"])
self.student = AutoModelForSequenceClassification.from_pretrained(hparams["student_model_name_or_path"])

and I want to save only student weights during checkpointing, how should I proceed?

Thank you

[tchaton]

Dear @ari9dam,

Here is one option to save only the student weights. This will drop in-place the parameters associated to the teacher from the checkpoint.

class MyModel(LightningModule)

    def on_save_checkpoint(self, checkpoint):
        keys = checkpoint["state_dict"].keys()
        for key in keys:
            if "teacher" in key:
                checkpoint["state_dict"].pop(key)

However, you will need to make sure to reload with strict=False mode.

Best,
T.C