1

I have a pre-trained model weight (as .pth) and it's configuration (as .yaml) and I want to fine-tune this model on my downstream task. The only problem is that I have 1 class while the pre trained model has 5 classes and when I have fine tuned my model with Detectron2, it gives me results for all the 5 classes instead of my 1 class. How can I deal with that scenario?

This is the exact tutorial which I am following but instead of training my classes on all 5 classes as thing_classes= ['None','text', 'title', 'list', 'table', 'figure'], I want to train just on one class as [text]. Author has answered but it did not help me as when I got the results during testing, I got results for all the 5 classes.

Pre-trained Model Weight Pre- trained Model Config

I have put 'category_id' of every instance as 0 (because I have just 1 class).

Below is the code where I have registered the data and everything and there is no problem with training, model trains well

from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor, DefaultTrainer

!wget -O ./faster_rcnn_R_50_FPN_3x.pth 'https://www.dropbox.com/s/dgy9c10wykk4lq4/model_final.pth?dl=1'

!wget -O ./faster_rcnn_R_50_FPN_3x.yaml 'https://www.dropbox.com/s/f3b12qc4hc0yh4m/config.yml?dl=1'

cfg = get_cfg()
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # Just one class predictions

cfg.merge_from_file("./faster_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS= './faster_rcnn_R_50_FPN_3x.pth' # layout parser Pre trained weights

cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.0025
cfg.SOLVER.MAX_ITER = 50 #adjust up if val mAP is still rising, adjust down if overfit
cfg.SOLVER.GAMMA = 0.05
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 4

cfg.DATASETS.TRAIN = (Data_Resister_training,)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()
Deshwal
  • 3,436
  • 4
  • 35
  • 94

1 Answers1

0

Not sure if this will fix it, but try inverting the merge_from_file call and the number of classes setting:

cfg = get_cfg()
cfg.merge_from_file("./faster_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # Just one class predictions
...

Maybe that parameter gets overwritten.

ClaudiaR
  • 3,108
  • 2
  • 13
  • 27
  • Okay. So you mean to say that my approach is correct and the configuration is getting overwritten somewhere? – Deshwal Aug 26 '22 at 04:00
  • Yes, in that place. Because you first set the number of classes and then merge the file. But I think it should be the other way around. Because if you open the file, you’ll see that inside the file is set `MODEL.ROI_HEADS.NUM_CLASSES`. Maybe, since it is set before in your script, before merging, after the merging your additional configuration is lost. I’m not entirely sure that this will fix it, but you should give it a try – ClaudiaR Aug 26 '22 at 05:17
  • Yes, I think so too. Thank you. Also, `cfg.MODEL.WEIGHTS` should be loaded before `cfg.MODEL.ROI_HEADS.NUM_CLASSES` or after? – Deshwal Aug 28 '22 at 10:19
  • I'm not sure this one makes any difference. To be sure I would add `cfg.MODEL.ROI_HEADS.NUM_CLASSES` after. – ClaudiaR Aug 28 '22 at 10:23
  • Yes there isn't any clear instructions on where. Also, I need to turn off the `cfg.MODEL.MASK_ON` to `False` as I'm not using `Mask RCNN` so it's position would make a difference? I think that has to be done after the weights as they have done the same in [faster RCNN](https://github.com/facebookresearch/detectron2/blob/main/configs/COCO-Detection/faster_rcnn_R_50_C4_1x.yaml) vs [Mask RCNN](https://github.com/facebookresearch/detectron2/blob/main/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_C4_1x.yaml) `yaml` files. – Deshwal Aug 28 '22 at 10:29
  • Yes, exactly, to configure my yaml I've followed as much as possible the files they had. – ClaudiaR Aug 28 '22 at 10:31