Write Models

If you are trying to do something completely new, you may wish to implement a model entirely from scratch within detectron2. However, in many situations you may be interested in modifying or extending some components of an existing model. Therefore, we also provide a registration mechanism that lets you override the behavior of certain internal components of standard models.

For example, to add a new backbone, import this code in your code:

from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec

@BACKBONE_REGISTRY.register()
class ToyBackBone(Backbone):
  def __init__(self, cfg, input_shape):
    # create your own backbone
    self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=16, padding=3)

  def forward(self, image):
    return {"conv1": self.conv1(image)}

  def output_shape(self):
    return {"conv1": ShapeSpec(channels=64, stride=16)}

Then, you can use cfg.MODEL.BACKBONE.NAME = 'ToyBackBone' in your config object. build_model(cfg) will then call your ToyBackBone instead.

As another example, to add new abilities to the ROI heads in the Generalized R-CNN meta-architecture, you can implement a new ROIHeads subclass and put it in the ROI_HEADS_REGISTRY. See densepose in detectron2 for an example that implements new ROIHeads.

Other registries can be found in API documentation. You can register components in these registries to customize different parts of a model, or the entire model.