Spaces:
Running
on
Zero
Running
on
Zero
| # Write Models | |
| If you are trying to do something completely new, you may wish to implement | |
| a model entirely from scratch. However, in many situations you may | |
| be interested in modifying or extending some components of an existing model. | |
| Therefore, we also provide mechanisms that let users override the | |
| behavior of certain internal components of standard models. | |
| ## Register New Components | |
| For common concepts that users often want to customize, such as "backbone feature extractor", "box head", | |
| we provide a registration mechanism for users to inject custom implementation that | |
| will be immediately available to use in config files. | |
| For example, to add a new backbone, import this code in your code: | |
| ```python | |
| from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec | |
| @BACKBONE_REGISTRY.register() | |
| class ToyBackbone(Backbone): | |
| def __init__(self, cfg, input_shape): | |
| super().__init__() | |
| # 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)} | |
| ``` | |
| In this code, we implement a new backbone following the interface of the | |
| [Backbone](../modules/modeling.html#detectron2.modeling.Backbone) class, | |
| and register it into the [BACKBONE_REGISTRY](../modules/modeling.html#detectron2.modeling.BACKBONE_REGISTRY) | |
| which requires subclasses of `Backbone`. | |
| After importing this code, detectron2 can link the name of the class to its implementation. Therefore you can write the following code: | |
| ```python | |
| cfg = ... # read a config | |
| cfg.MODEL.BACKBONE.NAME = 'ToyBackbone' # or set it in the config file | |
| model = build_model(cfg) # it will find `ToyBackbone` defined above | |
| ``` | |
| As another example, to add new abilities to the ROI heads in the Generalized R-CNN meta-architecture, | |
| you can implement a new | |
| [ROIHeads](../modules/modeling.html#detectron2.modeling.ROIHeads) subclass and put it in the `ROI_HEADS_REGISTRY`. | |
| [DensePose](../../projects/DensePose) | |
| and [MeshRCNN](https://github.com/facebookresearch/meshrcnn) | |
| are two examples that implement new ROIHeads to perform new tasks. | |
| And [projects/](../../projects/) | |
| contains more examples that implement different architectures. | |
| A complete list of registries can be found in [API documentation](../modules/modeling.html#model-registries). | |
| You can register components in these registries to customize different parts of a model, or the | |
| entire model. | |
| ## Construct Models with Explicit Arguments | |
| Registry is a bridge to connect names in config files to the actual code. | |
| They are meant to cover a few main components that users frequently need to replace. | |
| However, the capability of a text-based config file is sometimes limited and | |
| some deeper customization may be available only through writing code. | |
| Most model components in detectron2 have a clear `__init__` interface that documents | |
| what input arguments it needs. Calling them with custom arguments will give you a custom variant | |
| of the model. | |
| As an example, to use __custom loss function__ in the box head of a Faster R-CNN, we can do the following: | |
| 1. Losses are currently computed in [FastRCNNOutputLayers](../modules/modeling.html#detectron2.modeling.FastRCNNOutputLayers). | |
| We need to implement a variant or a subclass of it, with custom loss functions, named `MyRCNNOutput`. | |
| 2. Call `StandardROIHeads` with `box_predictor=MyRCNNOutput()` argument instead of the builtin `FastRCNNOutputLayers`. | |
| If all other arguments should stay unchanged, this can be easily achieved by using the [configurable `__init__`](../modules/config.html#detectron2.config.configurable) mechanism: | |
| ```python | |
| roi_heads = StandardROIHeads( | |
| cfg, backbone.output_shape(), | |
| box_predictor=MyRCNNOutput(...) | |
| ) | |
| ``` | |
| 3. (optional) If we want to enable this new model from a config file, registration is needed: | |
| ```python | |
| @ROI_HEADS_REGISTRY.register() | |
| class MyStandardROIHeads(StandardROIHeads): | |
| def __init__(self, cfg, input_shape): | |
| super().__init__(cfg, input_shape, | |
| box_predictor=MyRCNNOutput(...)) | |
| ``` | |