How to change number of classes for an already trained model(Retinanet)

Hey everyone,
I have trained my retiannet model setting classes=3, now I want to use the weights of the same model to train for another dataset that has number of classes=2.
Can anyone tell me how can i change the number classes after reloading the model?
I have tried forcing the 2 classes databunch into learner and loaded the model with torch.load .
The learner initializes without error , and can even perform inference for the new dataset, but when I do it gives out of index error.
Thanks, everyone for this great learning forum.

I’m not sure if it is possible. Usually, people use a pretrained model for the backbone, not the entire model. Because for Object Detection, you must change the dimension of all the cell-level predictions.

1 Like

Yeah exactly, but is there any workaround for it?
Because the data i have is not enough to train a retinanet model, so i use a dataset from the challenge which has 3 classes(including bg), how can I perform transfer learning for my new data which has only 2 classes?
If you know of a technique or patched code so that I can make the already existing model work for new data bunch it would be really helpful.
Thanks and regards,

Also, the 2 classes of the new dataset are the same as the first 2 classes of the challenge dataset.

Why not just extracts the 2 classes from your old dataset + your new data set, and train a new model?

1 Like

Yeah, I had the same idea of removing the last class from the old dataset, but I was doubtful of its impact on model learning as a hard negative class is also important for my model to differentiate between the positive class and background.
If I find no patch for changing the classifier head in the SSD I’ll resort to this method.
Thanks @dhoa