如何使用 Detectron2 微调自定义数据和类的对象检测模型？

Question

我有一个预先训练好的模型权重（.pth）和它的配置（.yaml），我想在我的下游任务中微调这个模型。唯一的问题是我有 1 个班级，而预训练模型有 5 个班级当我使用Detectron2 微调我的模型时，它会为我提供所有 5 个类的结果，而不是我的 1 个类。我该如何处理这种情况？

This is the exact tutorial which I am following 但我不想像thing_classes= ['None','text', 'title', 'list', 'table', 'figure'] 那样在所有 5 个课程上训练我的课程，我只想像 [text] 那样训练一个课程。作者已经回答了，但它对我没有帮助，因为当我在测试期间得到结果时，我得到了所有 5 个类的结果。

Pre-trained Model Weight Pre- trained Model Config

我已将每个实例的 'category_id' 设为 0（因为我只有一堂课）。

下面是我注册数据和所有内容的代码训练没有问题，模型训练得很好

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

不确定这是否会解决它，但尝试反转 merge_from_file 调用和类数设置：

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
...

也许该参数被覆盖。