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.
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.
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!