Load model weights into a new created model

1

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(
    encoder_name=ENCODER, 
    encoder_depth=ENCODER_DEPTH,
    encoder_weights=ENCODER_WEIGHTS, 
    decoder_channels=DECODER_CHANNELS,
    classes=len(CLASSES), 
    activation=ACTIVATION,
    decoder_attention_type=attention,
    #aux_params=dict(dropout=0,classes=len(CLASSES))
)

Then I loaded the weights:

model.load_state_dict(torch.load(path+"/model_loss.pth"))

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

2

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?

3

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

4

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?

5

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