You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

write-models.md 1.7 kB

3 years ago
123456789101112131415161718192021222324252627282930313233343536
  1. # Write Models
  2. If you are trying to do something completely new, you may wish to implement
  3. a model entirely from scratch within detectron2. However, in many situations you may
  4. be interested in modifying or extending some components of an existing model.
  5. Therefore, we also provide a registration mechanism that lets you override the
  6. behavior of certain internal components of standard models.
  7. For example, to add a new backbone, import this code in your code:
  8. ```python
  9. from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec
  10. @BACKBONE_REGISTRY.register()
  11. class ToyBackBone(Backbone):
  12. def __init__(self, cfg, input_shape):
  13. # create your own backbone
  14. self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=16, padding=3)
  15. def forward(self, image):
  16. return {"conv1": self.conv1(image)}
  17. def output_shape(self):
  18. return {"conv1": ShapeSpec(channels=64, stride=16)}
  19. ```
  20. Then, you can use `cfg.MODEL.BACKBONE.NAME = 'ToyBackBone'` in your config object.
  21. `build_model(cfg)` will then call your `ToyBackBone` instead.
  22. As another example, to add new abilities to the ROI heads in the Generalized R-CNN meta-architecture,
  23. you can implement a new
  24. [ROIHeads](../modules/modeling.html#detectron2.modeling.ROIHeads) subclass and put it in the `ROI_HEADS_REGISTRY`.
  25. See [densepose in detectron2](https://github.com/facebookresearch/detectron2/tree/master/projects/DensePose)
  26. for an example that implements new ROIHeads.
  27. Other registries can be found in [API documentation](../modules/modeling.html#model-registries).
  28. You can register components in these registries to customize different parts of a model, or the
  29. entire model.

No Description