在部分COCO数据集上训练Detectron2

时间:2020-07-08 14:16:17

标签: python machine-learning deep-learning pytorch detectron

我正在尝试使用Detectron2和COCO数据集训练模型以进行车辆和人员检测,并且在模型加载方面遇到问题。

我在SO和https://github.com/immersive-limit/coco-manager(filter.py文件)代码上使用过帖子,以过滤COCO数据集,使其仅包含来自“人员”,“汽车”,“自行车”,“卡车”类的注释和图像”和“自行车”。现在我的目录结构是:

main
  - annotations:
    - instances_train2017_filtered.json
    - instances_val2017_filtered.json
  - images:
    - train2017_filtered (lots of images inside)
    - val2017_filtered (lots of images inside)

基本上,我在这里所做的唯一一件事就是删除了与那些类不对应的文档和图像,并更改了它们的ID(因此它们从1到5)。

然后,我使用了Detectron2教程中的代码:

import random

import cv2
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.data.datasets import register_coco_instances
from detectron2.engine import DefaultTrainer, DefaultPredictor
from detectron2.config import get_cfg
import os

from detectron2.model_zoo import model_zoo
from detectron2.utils.visualizer import Visualizer

register_coco_instances("train",
                        {},
                        "/home/jakub/Projects/coco/annotations/instances_train2017_filtered.json",
                        "/home/jakub/Projects/coco/images/train2017_filtered/")

register_coco_instances("val",
                        {},
                        "/home/jakub/Projects/coco/annotations/instances_val2017_filtered.json",
                        "/home/jakub/Projects/coco/images/val2017_filtered/")

metadata = MetadataCatalog.get("train")
dataset_dicts = DatasetCatalog.get("train")

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 300
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 5

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")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.DATASETS.TEST = ("val", )
predictor = DefaultPredictor(cfg)

img = cv2.imread("demo/input.jpg")
outputs = predictor(img)

for d in random.sample(dataset_dicts, 1):
    im = cv2.imread(d["file_name"])
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=metadata,
                   scale=0.8)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2.imwrite('demo/output_retrained.jpg', out.get_image()[:, :, ::-1])

在培训期间,出现以下错误:

Unable to load 'roi_heads.box_predictor.cls_score.weight' to the model due to incompatible shapes: (81, 1024) in the checkpoint but (6, 1024) in the model!
Unable to load 'roi_heads.box_predictor.cls_score.bias' to the model due to incompatible shapes: (81,) in the checkpoint but (6,) in the model!
Unable to load 'roi_heads.box_predictor.bbox_pred.weight' to the model due to incompatible shapes: (320, 1024) in the checkpoint but (20, 1024) in the model!
Unable to load 'roi_heads.box_predictor.bbox_pred.bias' to the model due to incompatible shapes: (320,) in the checkpoint but (20,) in the model!
Unable to load 'roi_heads.mask_head.predictor.weight' to the model due to incompatible shapes: (80, 256, 1, 1) in the checkpoint but (5, 256, 1, 1) in the model!
Unable to load 'roi_heads.mask_head.predictor.bias' to the model due to incompatible shapes: (80,) in the checkpoint but (5,) in the model!

尽管减少了训练期间的total_loss,但是该模型无法预测训练后有用的任何东西。我知道由于大小不匹配(我减少了类的数量),我应该得到警告,这在互联网上是正常的,但是在每个错误行之后我都不会“跳过”。我认为该模型实际上未在此处加载任何内容,我想知道为什么以及如何解决此问题。

编辑

为进行比较,在几乎相同的情况下,类似的行为被报告为问题,但在每个错误行的末尾都“跳过了”,从而使它们有效地发出警告,而不是错误: https://github.com/facebookresearch/detectron2/issues/196

1 个答案:

答案 0 :(得分:1)

这个“警告”基本上是说您正在尝试从在不同数量的类上训练的模型初始化权重。正如您所阅读的那样,这是您所期望的。

我怀疑您没有从训练中获得任何结果,因为您的 MetadataCatalog 没有设置“thing_classes”属性。你只是打电话

MetadataCatalog.get("train")

打电话

MetadataCatalog.get("train").set(thing_classes=["person", "car", "bike", "truck", "bicycle"])

应该可以解决问题,但如果没有,我很确定您的 json 已损坏。