How to Use the Detectron2 Model Zoo (for Object Detection)

Detectron2 is Facebook's open source library for implementing state-of-the-art computer vision techniques in PyTorch. Facebook introduced Detectron2 in October 2019 as a complete rewrite of Detectron (which was implemented in Caffe). Facebook uses Detectron2 in a wide array of their products, including Portal, and notes the framework accelerates the feedback loop between research and production.

In releasing Detectron2, the Facebook Artificial Intelligence Research team also released a model zoo. The Detectron2 model zoo includes pre-trained models for a variety of tasks: object detection, semantic segmentation, and keypoint detection.

Using Detectron2 for Object Detection

The Roboflow team has published a Detectron2 tutorial on object detection, including a Detectron2 Colab notebook. (If you haven't yet followed that tutorial for training Detectron2, that is the recommended starting point.)

How to Train Detectron2 on Custom Object Detection Data

In this post, we will walk through how to train Detectron2 to detect customobjects in this Detectron2 Colab notebook[https://colab.research.google.com/drive/1-TNOcPm3Jr3fOJG8rnGT9gh60mHUsvaW#scrollTo=kc8MmgZugZWR]. After reading, you will be able to train your custom Detectron2 detector bychangi…

Roboflow BlogJacob Solawetz

Our prior Detectron2 tutorial assumes using a specific Faster-RCNN from the model zoo. However, the notebook can be customized to any other object detection model from the model library with a simple change.

In the notebook, when we load our model for training (line 10 below), note that we call for 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")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.001

cfg.SOLVER.WARMUP_ITERS = 1000
cfg.SOLVER.MAX_ITER = 3000 #adjust up if val mAP is still rising, adjust down if overfit
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 #your number of classes + 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()

However, if we visit the object detection model zoo in Detectron2, we see there are multiple implementations of Faster R-CNN available as well as RetinaNet and RetinaNet + Faster R-CNN.

To replace the YAML file with an alternative architecture (and pre-configured training checkpoint), simply:

1. Right click the model name in the lefthand column
2. Copy the link
3. Replace the link in the Colab notebook with the newly copied link

This new model YAML file then replaces the architecture, and training starts from that same pre-trained checkpoint. It's that easy.

Happy training!