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.