Detectron2 是由 Facebook 的人工智能研究团队在 2019 年 10 月推出的一个目标检测框架,它是 Detectron 的完全重写版本,后者是基于 Caffe 实现的。Detectron2 被广泛应用于 Facebook 的多种产品中,包括 Portal,并且该框架加速了研究与生产之间的反馈循环。
在发布 Detectron2 时,Facebook 人工智能研究团队还发布了一个模型库。Detectron2 模型库包含了多种任务的预训练模型:目标检测、语义分割和关键点检测。这些预训练模型为研究人员和开发者提供了一个强大的起点,以便他们能够快速地将模型应用到自己的数据集上。
Roboflow 团队发布了一个关于目标检测的 Detectron2 教程,包括一个 Detectron2 Colab 笔记本。如果还没有按照那个教程训练 Detectron2,那么这是推荐的起点。本文将带了解如何在 Detectron2 Colab 笔记本中训练 Detectron2 来检测自定义对象。
本文将详细介绍如何在Detectron2Colab 笔记本中训练 Detectron2 来检测自定义对象。阅读本文后,将能够通过简单的更改来训练自定义 Detectron2 检测器。
在之前的 Detectron2 教程中,假设使用的是模型库中的特定 Faster-RCNN 模型。然而,通过简单的更改,笔记本可以定制为模型库中的任何其他目标检测模型。
在笔记本中加载模型进行训练时(如下代码第 10 行所示),调用了 COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml:
from detectron2.config import get_cfg
import os
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = ("my_dataset_val",)
cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml") # 从模型库初始化训练
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.WARMUP_ITERS = 1000
cfg.SOLVER.MAX_ITER = 3000 # 如果验证 mAP 仍在上升,则上调;如果过拟合,则下调
cfg.SOLVER.STEPS = (1000, 1500)
cfg.SOLVER.GAMMA = 0.05
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 13 # 类别数量 + 1
cfg.TEST.EVAL_PERIOD = 500
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = CocoTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
使用 Faster R-CNN 在 Detectron2 中进行训练。然而,如果查看 Detectron2 的目标检测模型库,会发现有多种 Faster R-CNN 实现,以及 RetinaNet 和 RetinaNet + Faster R-CNN。
要替换 YAML 文件为另一种架构(以及预配置的训练检查点),只需:
这个新的模型 YAML 文件随后替换了架构,并且训练从同一个预训练检查点开始。就是这么简单。