Load model weights into a new created model

Hello, I trained a Unet Segmentation Model without Dropout in the SegmentationHead module.
I added the module to my network but I dont want to retrain it from scratch. I was wondering if I could initialize my model with dropout module on and then just load the weights from the model without dropout.

I did this:

model = smp.Unet(

Then I loaded the weights:


However, I am getting the following error:

‘Unet’ object has no attribute ‘copy’

Any idea if this is possible and if so, what am I doing wrong?

Kind regards

Are you able to load the state_dict without any modifications to the model?
The current error message seems to indicate that torch.load might load the complete model, not a state_dict.
Could this be the case, i.e. did you save the complete model instead of the state_dict?

Yes, I used torch.save() method to save the model at its best epoch.

I managed to make this work, I loaded the model and saved the dict and then I created a Unet with dropout in the architecture and loaded the dict.

However, I am getting this error:

RuntimeError: Error(s) in loading state_dict for Unet:
	Missing key(s) in state_dict: "segmentation_head.1.weight", "segmentation_head.1.bias". 
	Unexpected key(s) in state_dict: "segmentation_head.0.weight", "segmentation_head.0.bias".

I believe this is because of the modification I did in the SegmentatioHead by adding dropout. I want to load the unexpected keys as they are the missing keys. Is there any way to do this?

You could iterate the state_dict and change the name of these keys as shown here.