Loss Module with inner Network

[vasl12]

I have a loss module which is part of my lightning module with its own inner pretrained vgg network.
The problem comes when I am trying to use the checkpoint (that is saved automatically) to resume training or to test my model.
Then I get an error unexpected key(s) in state_dict pytorch lightning which points to the keys of the network which is part of the loss function.

Is there a way to load my model properly?

[tchaton]

Dear @vasl12,

You have multiple options:

• pass strict=False
• drop the key before saving the weights. Use on_save_checkpoint hook to access the checkpoint and drop the associated key.
• re-create the missing module in your model

Best,
T.C

[vasl12]

Dear @tchaton,

thank you very much for your answer. Regarding the first option, I am not sure in which function to pass it, since I am trying to load the checkpoint using the Trainer. This makes me also uncertain on how to implement the third option, since the code fails when I am trying to load the model, and in fact I have not changed the model itself.

I appreciate your help! Thank you very much!

[tchaton]

I think option 2 could work. Mind trying ?

class MyModel(LightningModule):

    def on_save_checkpoint(self, trainer, pl_moduel, checkpoint):
        state_dict = checkpoint["state_dict"]
        state_dict.pop("my_un_wanted_key")

[tchaton]

Here is option 1.

class MyModel(LightningModule):

    ...

model = MyModel()
state_dict = torch.load(...)
state_dict.pop("un_wanted_keys)
model.load_state_dict(state_dict)

or

class MyModel(LightningModule):

    ...


model = MyModel.load_from_checkpoint(path_to_checkpoint, strict=False)