Freezing Backbone and training the detection head using torchvision detection models

vision
1

I’m working with a COCO-style custom dataset and would like to freeze the backbone and train only the detection head for the following torchvision models pretrained on COCO:

  • fasterrcnn_resnet50_fpn
  • ssd300_vgg16
  • retinanet_resnet50_fpn
  • fcos_resnet50_fpn

My goal is to:

  1. Load each model with pretrained backbone weights,
  2. Replace or reinitialize the detection head to match my custom number of classes,
  3. Freeze the backbone (i.e., no gradient updates),
  4. Train only the detection head.
  5. Do inference and compare metrics like mAP etc

Can someone provide a reference implementation or point me in the right direction on how to do this cleanly?

Thanks!