Detectron2 source code reading 4-- registrar construction model

Preface

  This paper mainly introduces detectron2 How to build a model . This article will first introduce the Registrar , Then it introduces how to use the Registrar to register the model , Finally, let's introduce the construction process . If you are interested, you can have a look at mmdetection Medium register , You will find that the two excellent frameworks use the same design ideas .

1、Registry Introduce

  Registry You can understand it as a dictionary that can store classes . such as {‘BackBone’, resnet}. Look at the source ：

class Registry(Iterable[Tuple[str, Any]]):
    """ .. code-block:: python BACKBONE_REGISTRY = Registry('BACKBONE') To register an object: .. code-block:: python @BACKBONE_REGISTRY.register() class MyBackbone(): ... Or: .. code-block:: python BACKBONE_REGISTRY.register(MyBackbone) """

    def __init__(self, name: str) -> None:
        """ Args: name (str): the name of this registry """
        self._name: str = name
        self._obj_map: Dict[str, Any] = {
    }

    def _do_register(self, name: str, obj: Any) -> None:
        assert (
            name not in self._obj_map
        ), "An object named '{}' was already registered in '{}' registry!".format(
            name, self._name
        )
        self._obj_map[name] = obj

    def register(self, obj: Any = None) -> Any:
        """ Register the given object under the the name `obj.__name__`. Can be used as either a decorator or not. See docstring of this class for usage. """
        if obj is None:
            # used as a decorator
            def deco(func_or_class: Any) -> Any:
                name = func_or_class.__name__
                self._do_register(name, func_or_class)
                return func_or_class

            return deco

        # used as a function call
        name = obj.__name__
        self._do_register(name, obj)

    def get(self, name: str) -> Any:
        ret = self._obj_map.get(name)
        if ret is None:
            raise KeyError(
                "No object named '{}' found in '{}' registry!".format(name, self._name)
            )
        return ret

  Note describes two ways to register a new model .register Function will later be used as a decorator to register the model , and get The model will be built according to key To index the corresponding class . May not understand , Please jump ：mmdetection And Registry Introduce .

2、 structure ResNet50 For example

2.1. structure ResNet class

  Code from detectron2/modeling/backbone In the folder .

  Take a look first backbone.py, It defines a parent class , It's all backbone Parent class of , There is nothing to say here .

class Backbone(nn.Module, metaclass=ABCMeta):
    """ Abstract base class for network backbones. """

    def __init__(self):
        super().__init__()

    @abstractmethod
    def forward(self):
        pass

    @property
    def size_divisibility(self) -> int:
        return 0

    def output_shape(self):
        # this is a backward-compatible default
        return {
    
            name: ShapeSpec(
                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
            )
            for name in self._out_features
        }

  Let's see resnet.py. It mainly inherits from the above parent class , structure ResNet.

class ResNet(Backbone):
    """ Implement :paper:`ResNet`. """

    def __init__(self, stem, stages, num_classes=None, out_features=None):

2.2 utilize Registry register ResNet

  The registrar is used as follows ：

BACKBONE_REGISTRY = Registry("BACKBONE")    #  Instantiate a registrar 
@BACKBONE_REGISTRY.register()               #  With registrar register To decorate ResNet class 
def build_resnet_backbone(cfg, input_shape):

By , After the program runs , stay BACKBONE_REGISTRY The register contains ResNet class . And then through build.py Medium build_backbone The interface function , You can instantiate one ResNet.

def build_backbone(cfg, input_shape=None):
    """ Build a backbone from `cfg.MODEL.BACKBONE.NAME`. Returns: an instance of :class:`Backbone` """
    if input_shape is None:
        input_shape = ShapeSpec(channels=len(cfg.MODEL.PIXEL_MEAN))

    backbone_name = cfg.MODEL.BACKBONE.NAME
    backbone = BACKBONE_REGISTRY.get(backbone_name)(cfg, input_shape)  # .get Get class , and cfg Instantiate the class 
    assert isinstance(backbone, Backbone)
    return backbone                                    #  Build a resnet

3、SparseRCNN

  Now if you want to build your own model , How to use the Registrar ？ Here I take the open source project SparseRCNN For example , The source code uses detectron2 Build . Create a new one project Catalog , And then create a detector.py file . The content of this function is ：

from detectron2.modeling import META_ARCH_REGISTRY, build_backbone, detector_postprocess
__all__ = ["SparseRCNN"]

@META_ARCH_REGISTRY.register()
class SparseRCNN(nn.Module):

  Follow ResNet Build in the same way , First, import. modeling Medium META_ARCH_REGISTRY, The class is then SparseRCNN Sign up . After the detectron2/engine/defaults.py Call in build_model Just instantiate a network .

summary

 detectron2 With the help of the Registrar, it is very convenient to build the network .