Merge pull request #253 from hellock/registry

Use registry to manage modules

mmdet/models/necks/fpn.py

 import torch.nn as nn
 import torch.nn.functional as F
-from ..utils import ConvModule
 from mmcv.cnn import xavier_init
 
+from ..utils import ConvModule
+from ..registry import NECKS
+
 
+@NECKS.register_module
 class FPN(nn.Module):
 
     def __init__(self,

mmdet/models/registry.py

+import torch.nn as nn
+
+
+class Registry(object):
+
+    def __init__(self, name):
+        self._name = name
+        self._module_dict = dict()
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def module_dict(self):
+        return self._module_dict
+
+    def _register_module(self, module_class):
+        """Register a module.
+
+        Args:
+            module (:obj:`nn.Module`): Module to be registered.
+        """
+        if not issubclass(module_class, nn.Module):
+            raise TypeError(
+                'module must be a child of nn.Module, but got {}'.format(
+                    type(module_class)))
+        module_name = module_class.__name__
+        if module_name in self._module_dict:
+            raise KeyError('{} is already registered in {}'.format(
+                module_name, self.name))
+        self._module_dict[module_name] = module_class
+
+    def register_module(self, cls):
+        self._register_module(cls)
+        return cls
+
+
+BACKBONES = Registry('backbone')
+NECKS = Registry('neck')
+ROI_EXTRACTORS = Registry('roi_extractor')
+HEADS = Registry('head')
+DETECTORS = Registry('detector')

mmdet/models/roi_extractors/single_level.py

@@ -4,8 +4,10 @@ import torch
 import torch.nn as nn
 
 from mmdet import ops
+from ..registry import ROI_EXTRACTORS
 
 
+@ROI_EXTRACTORS.register_module
 class SingleRoIExtractor(nn.Module):
     """Extract RoI features from a single level feature map.