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

Question

我有一个预先训练的模型权重(作为.pth(和它的配置(作为.yaml(，我想在我的下游任务中微调这个模型。唯一的问题是我有1个类，而预训练的模型有5个类，当我用Detectron2微调模型时，它会给我所有5个类的结果，而不是我的1个类。我该如何应对这种情况？

这正是我正在学习的教程，但我不想把我的课全部训练成thing_classes= ['None','text', 'title', 'list', 'table', 'figure']，而是只想把一个课训练成[text]。作者已经回答了，但这对我没有帮助，因为当我在测试中得到结果时，我得到了所有5个班的结果。

预先训练的模型重量预先训练的模型配置

我已经将每个实例的'category_id'设置为0(因为我只有1个类(。

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

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

也许那个参数被覆盖了。