Detectron2 自定义目标检测训练指南

Detectron2 是由 Facebook 的人工智能研究团队在 2019 年 10 月推出的一个目标检测框架,它是 Detectron 的完全重写版本,后者是基于 Caffe 实现的。Detectron2 被广泛应用于 Facebook 的多种产品中,包括 Portal,并且该框架加速了研究与生产之间的反馈循环。

Detectron2模型库

在发布 Detectron2 时,Facebook 人工智能研究团队还发布了一个模型库。Detectron2 模型库包含了多种任务的预训练模型:目标检测、语义分割和关键点检测。这些预训练模型为研究人员和开发者提供了一个强大的起点,以便他们能够快速地将模型应用到自己的数据集上。

使用Detectron2进行目标检测

Roboflow 团队发布了一个关于目标检测的 Detectron2 教程,包括一个 Detectron2 Colab 笔记本。如果还没有按照那个教程训练 Detectron2,那么这是推荐的起点。本文将带了解如何在 Detectron2 Colab 笔记本中训练 Detectron2 来检测自定义对象。

在自定义数据上训练 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 文件为另一种架构(以及预配置的训练检查点),只需:

  • 右键单击左侧列中的模型名称
  • 复制链接
  • 将 Colab 笔记本中的链接替换为新复制的链接

这个新的模型 YAML 文件随后替换了架构,并且训练从同一个预训练检查点开始。就是这么简单。

沪ICP备2024098111号-1
上海秋旦网络科技中心:上海市奉贤区金大公路8218号1幢 联系电话:17898875485