Skip to content

Commit 375605f

Browse files
authored
add model registry (#760)
* add model registry * fixed infer scoep * fixed build func * add docstring * add md * support multi level * clean comments * add docs * fixed parent * add more doc * add value error, add docstring * fixed docs * change to local/global search * resolve comments * fixed test * update some docstring * update docs (minior) * update docs * update docs
1 parent 47825b1 commit 375605f

File tree

6 files changed

+429
-72
lines changed

6 files changed

+429
-72
lines changed

docs/registry.md

+98-20
Original file line numberDiff line numberDiff line change
@@ -9,40 +9,29 @@ In MMCV, registry can be regarded as a mapping that maps a class to a string.
99
These classes contained by a single registry usually have similar APIs but implement different algorithms or support different datasets.
1010
With the registry, users can find and instantiate the class through its corresponding string, and use the instantiated module as they want.
1111
One typical example is the config systems in most OpenMMLab projects, which use the registry to create hooks, runners, models, and datasets, through configs.
12+
The API reference could be find [here](https://mmcv.readthedocs.io/en/latest/api.html?highlight=registry#mmcv.utils.Registry).
1213

1314
To manage your modules in the codebase by `Registry`, there are three steps as below.
1415

15-
1. Create an registry
16-
2. Create a build method
17-
3. Use this registry to manage the modules
16+
1. Create a build method (optional, in most cases you can just use the default one).
17+
2. Create a registry.
18+
3. Use this registry to manage the modules.
19+
20+
`build_func` argument of `Registry` is to customize how to instantiate the class instance, the default one is `build_from_cfg` implemented [here](https://mmcv.readthedocs.io/en/latest/api.html?highlight=registry#mmcv.utils.build_from_cfg).
1821

1922
### A Simple Example
2023

2124
Here we show a simple example of using registry to manage modules in a package.
2225
You can find more practical examples in OpenMMLab projects.
2326

2427
Assuming we want to implement a series of Dataset Converter for converting different formats of data to the expected data format.
25-
We create directory as a package named `converters`.
28+
We create a directory as a package named `converters`.
2629
In the package, we first create a file to implement builders, named `converters/builder.py`, as below
2730

2831
```python
2932
from mmcv.utils import Registry
30-
3133
# create a registry for converters
3234
CONVERTERS = Registry('converter')
33-
34-
35-
# create a build function
36-
def build_converter(cfg, *args, **kwargs):
37-
cfg_ = cfg.copy()
38-
converter_type = cfg_.pop('type')
39-
if converter_type not in CONVERTERS:
40-
raise KeyError(f'Unrecognized task type {converter_type}')
41-
else:
42-
converter_cls = CONVERTERS.get(converter_type)
43-
44-
converter = converter_cls(*args, **kwargs, **cfg_)
45-
return converter
4635
```
4736

4837
Then we can implement different converters in the package. For example, implement `Converter1` in `converters/converter1.py`
@@ -51,7 +40,6 @@ Then we can implement different converters in the package. For example, implemen
5140

5241
from .builder import CONVERTERS
5342

54-
5543
# use the registry to manage the module
5644
@CONVERTERS.register_module()
5745
class Converter1(object):
@@ -71,5 +59,95 @@ If the module is successfully registered, you can use this converter through con
7159

7260
```python
7361
converter_cfg = dict(type='Converter1', a=a_value, b=b_value)
74-
converter = build_converter(converter_cfg)
62+
converter = CONVERTERS.build(converter_cfg)
63+
```
64+
65+
## Customize Build Function
66+
67+
Suppose we would like to customize how `converters` are built, we could implement a customized `build_func` and pass it into the registry.
68+
69+
```python
70+
from mmcv.utils import Registry
71+
72+
# create a build function
73+
def build_converter(cfg, registry, *args, **kwargs):
74+
cfg_ = cfg.copy()
75+
converter_type = cfg_.pop('type')
76+
if converter_type not in registry:
77+
raise KeyError(f'Unrecognized converter type {converter_type}')
78+
else:
79+
converter_cls = registry.get(converter_type)
80+
81+
converter = converter_cls(*args, **kwargs, **cfg_)
82+
return converter
83+
84+
# create a registry for converters and pass ``build_converter`` function
85+
CONVERTERS = Registry('converter', build_func=build_converter)
7586
```
87+
88+
Note: in this example, we demonstrate how to use the `build_func` argument to customize the way to build a class instance.
89+
The functionality is similar to the default `build_from_cfg`. In most cases, default one would be sufficient.
90+
`build_model_from_cfg` is also implemented to build PyTorch module in `nn.Sequentail`, you may directly use them instead of implementing by yourself.
91+
92+
## Hierarchy Registry
93+
94+
You could also build modules from more than one OpenMMLab frameworks, e.g. you could use all backbones in [MMClassification](https://github.com/open-mmlab/mmclassification) for object detectors in [MMDetection](https://github.com/open-mmlab/mmdetection), you may also combine an object detection model in [MMDetection](https://github.com/open-mmlab/mmdetection) and semantic segmentation model in [MMSegmentation](https://github.com/open-mmlab/mmsegmentation).
95+
96+
All `MODELS` registries of downstream codebases are children registries of MMCV's `MODELS` registry.
97+
Basically, there are two ways to build a module from child or sibling registries.
98+
99+
1. Build from children registries.
100+
101+
For example:
102+
103+
In MMDetection we define:
104+
105+
```python
106+
from mmcv.utils import Registry
107+
from mmcv.cnn import MODELS as MMCV_MODELS
108+
MODELS = Registry('model', parent=MMCV_MODELS)
109+
110+
@MODELS.register_module()
111+
class NetA(nn.Module):
112+
def forward(self, x):
113+
return x
114+
```
115+
116+
In MMClassification we define:
117+
118+
```python
119+
from mmcv.utils import Registry
120+
from mmcv.cnn import MODELS as MMCV_MODELS
121+
MODELS = Registry('model', parent=MMCV_MODELS)
122+
123+
@MODELS.register_module()
124+
class NetB(nn.Module):
125+
def forward(self, x):
126+
return x + 1
127+
```
128+
129+
We could build two net in either MMDetection or MMClassification by:
130+
131+
```python
132+
from mmdet.models import MODELS
133+
net_a = MODELS.build(cfg=dict(type='NetA'))
134+
net_b = MODELS.build(cfg=dict(type='mmcls.NetB'))
135+
```
136+
137+
or
138+
139+
```python
140+
from mmcls.models import MODELS
141+
net_a = MODELS.build(cfg=dict(type='mmdet.NetA'))
142+
net_b = MODELS.build(cfg=dict(type='NetB'))
143+
```
144+
145+
2. Build from parent registry.
146+
147+
The shared `MODELS` registry in MMCV is the parent registry for all downstream codebases (root registry):
148+
149+
```python
150+
from mmcv.cnn import MODELS as MMCV_MODELS
151+
net_a = MMCV_MODELS.build(cfg=dict(type='mmdet.NetA'))
152+
net_b = MMCV_MODELS.build(cfg=dict(type='mmcls.NetB'))
153+
```

mmcv/cnn/__init__.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
build_activation_layer, build_conv_layer,
1212
build_norm_layer, build_padding_layer, build_plugin_layer,
1313
build_upsample_layer, conv_ws_2d, is_norm)
14+
from .builder import MODELS, build_model_from_cfg
1415
# yapf: enable
1516
from .resnet import ResNet, make_res_layer
1617
from .utils import (INITIALIZERS, Caffe2XavierInit, ConstantInit, KaimingInit,
@@ -34,5 +35,5 @@
3435
'Linear', 'Conv2d', 'ConvTranspose2d', 'MaxPool2d', 'ConvTranspose3d',
3536
'MaxPool3d', 'Conv3d', 'initialize', 'INITIALIZERS', 'ConstantInit',
3637
'XavierInit', 'NormalInit', 'UniformInit', 'KaimingInit', 'PretrainedInit',
37-
'Caffe2XavierInit'
38+
'Caffe2XavierInit', 'MODELS', 'build_model_from_cfg'
3839
]

mmcv/cnn/builder.py

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import torch.nn as nn
2+
3+
from ..utils import Registry, build_from_cfg
4+
5+
6+
def build_model_from_cfg(cfg, registry, default_args=None):
7+
"""Build a PyTorch model from config dict(s). Different from
8+
``build_from_cfg``, if cfg is a list, a ``nn.Sequential`` will be built.
9+
10+
Args:
11+
cfg (dict, list[dict]): The config of modules, is is either a config
12+
dict or a list of config dicts. If cfg is a list, a
13+
the built modules will be wrapped with ``nn.Sequential``.
14+
registry (:obj:`Registry`): A registry the module belongs to.
15+
default_args (dict, optional): Default arguments to build the module.
16+
Defaults to None.
17+
18+
Returns:
19+
nn.Module: A built nn module.
20+
"""
21+
if isinstance(cfg, list):
22+
modules = [
23+
build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg
24+
]
25+
return nn.Sequential(*modules)
26+
else:
27+
return build_from_cfg(cfg, registry, default_args)
28+
29+
30+
MODELS = Registry('model', build_func=build_model_from_cfg)

0 commit comments

Comments
 (0)