Hi! I have a pretrained ssdlite320_mobilenet_v3_large model with a classification/regression head.
The model performs object detection + classification among several tasks. However, now I want to finetune the model so that it works on a reduced number of classes.
How should I go about retraining/finetuning the head part only? Is retraining even required?
I am trying something like this, but the output dimensions during training seem to be off.
model_old = SSDMobilenet(num_classes=20) #
model_old.load_state_dict(model_path, map_location=self.device)
# New model with user-defined number of classes
model = SSDMobilenet(num_classes=4)
# Use old model backbone with new classifier
model.backbone = model_old.torchvision_model.backbone
for param in model.parameters():
param.requires_grad = False
for param in model.torchvision_model.head.classification_head.parameters():
param.requires_grad = True