如何使用 Detectron2 为自定义数据和类微调 object 检测 model？

Question

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()

Answer 1

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.