Change the head of a loaded fastai model

So I have trained a resnet50(pretrained model) on a new Dataset A. For dataset A I did a categorical image classification with 4 categories and saved the new model weights using model.save(“modelAtrained”) .
Then I tried to load the weights of “modelAtrained” and train on a new dataset B using again a resnet50(pretrained model).

The code used is shown below:

#load the data of dataset B
data = ImageDataBunch.from_folder(path=“datasetB”, train=‘train’,
valid=‘val’,ds_tfms=get_transforms(), size=224, bs=batch_size[i])#, check_ext=False)
#Normalizing data based on Image net parameters
data.normalize(imagenet_stats)

#LOAD THE TRANSFER LEARNING MODEL
trans_model= cnn_learner(data, model, metrics=[accuracy])
#load the weights of previous model A
trans_model.load(“modelAtrained.pth”)

end of code

However an ERROR occurs. This is because the new Dataset B has 6 categories. So when the weights are loaded there is a mismatch between the head of the new model and the previous model A.
The head of the previous model A, has 4 outputs, however for this model 6 outputs are required.
I was wondering how could i change the head to 6 outputs in order to train the new dataset B.

Thank you and any feedback is much appreciated!

1 Like

This has been asked a lot on the forum. Look for something like transfer learning twice

1 Like

Your comment was really helpful! However even though I implemented the line and changed my code I ended up with another error of size mismatch.

Hi @cantonioupao - would need to see more of your code but I can point out one immediate issue:

In your code snippet (learn summary) you have specified the model to output 5 features vs in your first post you indicated you wanted it to do 6?

You may have accidentally copied the exact line from that thread and forgotten to update it to what you need?
ie -> current line:

learn.model[-1][-1]=nn.Linear(in_features=512,out_features=**5**, bias=True)   

you would want:

learn.model[-1][-1]=nn.Linear(in_features=512,out_features=**6**, bias=True) 

without the ** of course :slight_smile:

Also, you may want/need to restart your kernel and reload everything after making changes like the above just to make sure you didn’t hit a memory reload issue.

Otherwise, it’s basically saying the tensor sizes don’t match so the optimizer can’t adjust the weights properly for the update step. As noted a kernel restart might help ensure you have everything you want loaded, though at a minimum you want to update the model to output the category size you need.

Hope that is helpful!

3 Likes

Thank you very much for your reply! The issue is solved now. I didn’t unfreeze the model before applying the change to the header and that’s why i was getting this error. Thank you both for your help and fast response!:smiley::smiley::smiley::smiley:

3 Likes

@cantonioupao how do you freeze the layer and then change the output classes? I am a bit confused…
Please help