Transfer learning - freezing weights

[Nerolf05]

Hi together,
first thank you for sharing this amazing project.

1. What inputs you will provide, if any:

Input 1: COCO Instance Segmentation Baselines with Mask R-CNN ("COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml") with corresponding weights
Input 2: New custom dataset of 'building' class (satellite aerial images)

2. Expected Outcome:

I would like to apply transfer learning as suggested by https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html, in especially i would like to:
I: Freeze the model - except roi_heads - and train only the RPN head for some epochs
II: Afterwards unfreeze the model and fine-tune the whole model for some additional epochs with a reduced lr

Now I would like to ask, is there a way in detectron2 to freeze certain parts of a model or should the PyTorch API be used, like:

for param in model_conv.parameters():
    param.requires_grad = False

Is there an efficient way to freeze the whole model except the roi_heads as they cannot be loaded due to different number of classes?

I have also noticed, that some layers are "frozen" per default for the specified baseline

[idx for idx, x in enumerate(model.parameters()) if x.requires_grad == False]
-> [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26]

Should these layers always remain frozen, even in the case of fine-tuning (Phase II)?

Thanks !

[VGrondin]

Hi did you get it to work? Trying to do the same

[Nerolf05]

Hi @VGrondin,
at the end I used a 3 phase training for the MaskRCNN architecture, similar to a paper I have read:

def schedule_transfer_learning():
# Freeze whole model
for param in model.parameters():
    param.requires_grad = False
# Phase 1: Unfreeze only the roi_heads
for param in model.roi_heads.parameters():
    param.requires_grad = True
yield 1
# Phase 2: Unfreeze region proposal generator with reduced lr
for param in model.proposal_generator.parameters():
     param.requires_grad = True
yield 2
# Phase 3: Fine tuning of whole model on all layers with reduced lr
for param in model.parameters():
    param.requires_grad = True
return 3

Note that this scheduler should be called at the same time as the learning rate scheduler

[VGrondin]

Thank you your code helped me a lot. On top of that, I used this line of code to check the name of all parameters we can freeze/unfreeze:

for param in trainer.model.named_parameters():
        print(param[0])

[Kevin-Delnoije]

This technique looks really cool, where do you apply this code in plain_train_net.py or is it possible to implement when subclassing DefaultTrainer?

[Nerolf05]

I implemented it in plain_train_net.py in the training loop:

transfer_learning = schedule_transfer_learning(model)
next(transfer_learning)  # Initialize phase 1

# Actual training loop
with EventStorage(start_iter) as storage:
    for data, iteration in tqdm(zip(data_loader, range(start_iter, max_iter))):
        storage.iter = iteration
        loss_dict = model(data)
        losses = sum(loss_dict.values())
        assert torch.isfinite(losses).all(), loss_dict
        ...
        
       # Schedule transfer-learning
        if cfg.SOLVER.STEPS[steps_counter] == iteration:
            try:
                next(transfer_learning_gen)
            except StopIteration:
                continue
            finally:
                steps_counter = steps_counter +1 if (steps_counter +1) < len(cfg.SOLVER.STEPS) else steps_counter

But I assume this quick approach could also be added as a hook to DefaultTrainer similar to a learning-rate scheduler.

[ershashank42]

Hi here,
I am very new to this transfer learning. Didn't get much about what you guys discussed. Can someone please make it a little easy for me, how to re-train a model using transfer learning?
I'll share my code and the steps I did till now -

1. I have a custom dataset for 10000 images on curtains as a zip in COCO format.
2. I trained a model using custom dataset for curtains(curtain_train)
3. I ran trainer.train() on it. And got the model_final.pth file out of it.
   `for d in ["train", "valid"]:
   DatasetCatalog.register("curtain_" + d, lambda d=d: getCurtainsTrained(imgPath + d, jsonPath))
   MetadataCatalog.get("curtain_" + d).set(thing_classes=["curtain"])
   curtain_metadata = MetadataCatalog.get("curtain_train")

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("curtain_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml") # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 500
cfg.SOLVER.STEPS = []
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") # path to the model we just trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9 # set a custom testing threshold
predictor = DefaultPredictor(cfg)
`

Now I have another dataset with 5000 images on chairs. I don't want to re-register and re-train my model for 2 categories from scratch. I want to use the concept of transfer learning here. Please someone guide me by giving me sample code snippet by highlighting what to change and where to change in my code for the same.

It'll be really grateful of you if someone could reply on this asap. I'm working on a live project and it has a deadline for the same.
Thanks.

[beez2022]

It seems that we can set the following to freeze the weights of earlier layers? Hope someone more experienced can verify. Thank you:

cfg = get_cfg()
......
.....
trainer = DefaultTrainer(cfg)

to freeze the backbone and proposal_generator layers:

freeze = ['backbone', 'proposal_generator']
for k, v in trainer.model.named_parameters():
v.requires_grad = True
if any(x in k for x in freeze):
print(k)
v.requires_grad = False

trainer.resume_or_load(resume=False)
trainer.train()