diff --git a/docs/en/advanced_tutorials/registry.md b/docs/en/advanced_tutorials/registry.md index e121cd89cd..eb875961de 100644 --- a/docs/en/advanced_tutorials/registry.md +++ b/docs/en/advanced_tutorials/registry.md @@ -1,3 +1,302 @@ # Registry -Coming soon. Please refer to [chinese documentation](https://mmengine.readthedocs.io/zh_CN/latest/advanced_tutorials/registry.html). +OpenMMLab supports a rich collection of algorithms and datasets, therefore, many modules with similar functionality are implemented. For example, the implementations of `ResNet` and `SE-ResNet` are based on the classes `ResNet` and `SEResNet`, respectively, which have similar functions and interfaces and belong to the model components of the algorithm library. To manage these functionally similar modules, MMEngine implements the [registry](mmengine.registry.registry). Most of the algorithm libraries in OpenMMLab use `registry` to manage their modules, including [MMDetection](https://github.com/open-mmlab/mmdetection), [MMDetection3D](https://github.com/open-mmlab/mmdetection3d), [MMClassification](https://github.com/open-mmlab/mmclassification) and [MMEditing](https://github.com/open-mmlab/mmediting), etc. + +## What is a registry + +The [registry](mmengine.registry.Registry) in MMEngine can be considered as a union of a mapping table and a build function of modules. The mapping table maintains a mapping from strings to **classes or functions**, allowing the user to find the corresponding class or function with its name/notation. For example, the mapping from the string `"ResNet"` to the `ResNet` class. The module build function defines how to find the corresponding class or function based on a string and how to instantiate the class or call the function. For example, finding `nn.BatchNorm2d` and instantiating the `BatchNorm2d` module by the string `"bn"`, or finding the `build_batchnorm2d` function by the string `"build_batchnorm2d"` and then returning the result. The registries in MMEngine use the [build_from_cfg](mmengine.registry.build_from_cfg) function by default to find and instantiate the class or function corresponding to the string. + +The classes or functions managed by a registry usually have similar interfaces and functionality, so the registry can be treated as an abstraction of those classes or functions. For example, the registry `MODELS` can be treated as an abstraction of all models, which manages classes such as `ResNet`, `SEResNet` and `RegNetX` and constructors such as `build_ResNet`, `build_SEResNet` and `build_RegNetX`. + +## Getting started + +There are three steps required to use the registry to manage modules in the codebase. + +1. Create a registry. +2. Create a build method for instantiating the class (optional because in most cases you can just use the default method). +3. Add the module to the registry + +Suppose we want to implement a series of activation modules and want to be able to switch to different modules by just modifying the configuration without modifying the code. + +Let's create a regitry first. + +```python +from mmengine import Registry +# scope represents the domain of the registry. If not set, the default value is the package name. +# e.g. in mmdetection, the scope is mmdet +ACTIVATION = Registry('activation', scope='mmengine') +``` + +Then we can implement different activation modules, such as `Sigmoid`, `ReLU`, and `Softmax`. + +```python +import torch.nn as nn + +# use the register_module +@ACTIVATION.register_module() +class Sigmoid(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + print('call Sigmoid.forward') + return x + +@ACTIVATION.register_module() +class ReLU(nn.Module): + def __init__(self, inplace=False): + super().__init__() + + def forward(self, x): + print('call ReLU.forward') + return x + +@ACTIVATION.register_module() +class Softmax(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + print('call Softmax.forward') + return x +``` + +The key of using the registry module is to register the implemented modules into the `ACTIVATION` registry. With the `@ACTIVATION.register_module()` decorator added before the implemented module, the mapping between strings and classes or functions can be built and maintained by `ACTIVATION`. We can achieve the same functionality with `ACTIVATION.register_module(module=ReLU)` as well. + +By registering, we can create a mapping between strings and classes or functions via `ACTIVATION`. + +```python +print(ACTIVATION.module_dict) +# { +# 'Sigmoid': __main__.Sigmoid, +# 'ReLU': __main__.ReLU, +# 'Softmax': __main__.Softmax +# } +``` + +```{note} +The registry mechanism will only be triggered when the corresponded module file is imported, so we need to import the file somewhere or dynamically import the module using the ``custom_imports`` field to trigger the mechanism. Please refer to [Importing custom Python modules](config.md#import-the-custom-module) for more details. +``` + +Once the implemented module is successfully registered, we can use the activation module in the configuration file. + +```python +import torch + +input = torch.randn(2) + +act_cfg = dict(type='Sigmoid') +activation = ACTIVATION.build(act_cfg) +output = activation(input) +# call Sigmoid.forward +print(output) +``` + +We can switch to `ReLU` by just changing this configuration. + +```python +act_cfg = dict(type='ReLU', inplace=True) +activation = ACTIVATION.build(act_cfg) +output = activation(input) +# call ReLU.forward +print(output) +``` + +If we want to check the type of input parameters (or any other operations) before creating an instance, we can implement a build method and pass it to the registry to implement a custom build process. + +Create a `build_activation` function. + +```python +def build_activation(cfg, registry, *args, **kwargs): + cfg_ = cfg.copy() + act_type = cfg_.pop('type') + print(f'build activation: {act_type}') + act_cls = registry.get(act_type) + act = act_cls(*args, **kwargs, **cfg_) + return act +``` + +Pass the `buid_activation` to `build_func`. + +```python +ACTIVATION = Registry('activation', build_func=build_activation, scope='mmengine') + +@ACTIVATION.register_module() +class Tanh(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + print('call Tanh.forward') + return x + +act_cfg = dict(type='Tanh') +activation = ACTIVATION.build(act_cfg) +output = activation(input) +# build activation: Tanh +# call Tanh.forward +print(output) +``` + +```{note} +In the above example, we demonstrate how to customize the method of building an instance of a class using the `build_func`. +This is similar to the default `build_from_cfg` method. In most cases, using the default method will be fine. +``` + +MMEngine's registry can register classes as well as functions. + +```python +FUNCTION = Registry('function', scope='mmengine') + +@FUNCTION.register_module() +def print_args(**kwargs): + print(kwargs) + +func_cfg = dict(type='print_args', a=1, b=2) +func_res = FUNCTION.build(func_cfg) +``` + +## Advanced usage + +The registry in MMEngine supports hierarchical registration, which enables cross-project calls, meaning that modules from one project can be used in another project. Though there are other ways to implement this, the registry provides a much easier solution. + +To easily make cross-library calls, MMEngine provides twenty root registries, including: + +- RUNNERS: the registry for Runner. +- RUNNER_CONSTRUCTORS: the constructors for Runner. +- LOOPS: manages training, validation and testing processes, such as `EpochBasedTrainLoop`. +- HOOKS: the hooks, such as `CheckpointHook`, and `ParamSchedulerHook`. +- DATASETS: the datasets. +- DATA_SAMPLERS: `Sampler` of `DataLoader`, used to sample the data. +- TRANSFORMS: various data preprocessing methods, such as `Resize`, and `Reshape`. +- MODELS: various modules of the model. +- MODEL_WRAPPERS: model wrappers for parallelizing distributed data, such as `MMDistributedDataParallel`. +- WEIGHT_INITIALIZERS: the tools for weight initialization. +- OPTIMIZERS: registers all `Optimizers` and custom `Optimizers` in PyTorch. +- OPTIM_WRAPPER: the wrapper for Optimizer-related operations such as `OptimWrapper`, and `AmpOptimWrapper`. +- OPTIM_WRAPPER_CONSTRUCTORS: the constructors for optimizer wrappers. +- PARAM_SCHEDULERS: various parameter schedulers, such as `MultiStepLR`. +- METRICS: the evaluation metrics for computing model accuracy, such as `Accuracy`. +- EVALUATOR: one or more evaluation metrics used to calculate the model accuracy. +- TASK_UTILS: the task-intensive components, such as `AnchorGenerator`, and `BboxCoder`. +- VISUALIZERS: the management drawing module that draws prediction boxes on images, such as `DetVisualizer`. +- VISBACKENDS: the backend for storing training logs, such as `LocalVisBackend`, and `TensorboardVisBackend`. +- LOG_PROCESSORS: controls the log statistics window and statistics methods, by default we use `LogProcessor`. You may customize `LogProcessor` if you have special needs. + +### Use the module of the parent node + +Let's define a `RReLU` module in `MMEngine` and register it to the `MODELS` root registry. + +```python +import torch.nn as nn +from mmengine import Registry, MODELS + +@MODELS.register_module() +class RReLU(nn.Module): + def __init__(self, lower=0.125, upper=0.333, inplace=False): + super().__init__() + + def forward(self, x): + print('call RReLU.forward') + return x +``` + +Now suppose there is a project called `MMAlpha`, which also defines a `MODELS` and sets its parent node to the `MODELS` of `MMEngine`, which creates a hierarchical structure. + +```python +from mmengine import Registry, MODELS as MMENGINE_MODELS + +MODELS = Registry('model', parent=MMENGINE_MODELS, scope='mmalpha') +``` + +The following figure shows the hierarchy of `MMEngine` and `MMAlpha`. + +
+ +
+ +The [count_registered_modules](mmengine.registry.count_registered_modules) function can be used to print the modules that have been registered to MMEngine and their hierarchy. + +```python +from mmengine.registry import count_registered_modules + +count_registered_modules() +``` + +We define a customized `LogSoftmax` module in `MMAlpha` and register it to the `MODELS` in `MMAlpha`. + +```python +@MODELS.register_module() +class LogSoftmax(nn.Module): + def __init__(self, dim=None): + super().__init__() + + def forward(self, x): + print('call LogSoftmax.forward') + return x +``` + +Here we use the `LogSoftmax` in the configuration of `MMAlpha`. + +```python +model = MODELS.build(cfg=dict(type='LogSoftmax')) +``` + +We can also use the modules of the parent node `MMEngine` here in the `MMAlpha`. + +```python +model = MODELS.build(cfg=dict(type='RReLU', lower=0.2)) +# scope is optional +model = MODELS.build(cfg=dict(type='mmengine.RReLU')) +``` + +If no prefix is added, the `build` method will first find out if the module exists in the current node and return it if there is one. Otherwise, it will continue to look up the parent nodes or even the ancestor node until it finds the module. If the same module exists in both the current node and the parent nodes, we need to specify the `scope` prefix to indicate that we want to use the module of the parent nodes. + +```python +import torch + +input = torch.randn(2) +output = model(input) +# call RReLU.forward +print(output) +``` + +### Use the module of a sibling node + +In addition to using the module of the parent nodes, users can also call the module of a sibling node. + +Suppose there is another project called `MMBeta`, which, like `MMAlpha`, defines `MODELS` and set its parent node to `MMEngine`. + +```python +from mmengine import Registry, MODELS as MMENGINE_MODELS + +MODELS = Registry('model', parent=MMENGINE_MODELS, scope='mmbeta') +``` + +The following figure shows the registry structure of `MMAlpha` and `MMBeta`. + +
+ +
+ +Now we call the modules of `MMAlpha` in `MMBeta`. + +```python +model = MODELS.build(cfg=dict(type='mmalpha.LogSoftmax')) +output = model(input) +# call LogSoftmax.forward +print(output) +``` + +Calling a module of a sibling node requires the `scope` prefix to be specified in `type`, so the above configuration requires the prefix `mmalpha`. + +However, if you need to call several modules of a sibling node, each with a prefix, this requires a lot of modification. Therefore, `MMEngine` introduces the [DefaultScope](mmengine.registry.DefaultScope), with which `Registry` can easily support temporary switching of the current node to the specified node. + +If you need to switch the current node to the specified node temporarily, just set `_scope_` to the scope of the specified node in `cfg`. + +```python +model = MODELS.build(cfg=dict(type='LogSoftmax', _scope_='mmalpha')) +output = model(input) +# call LogSoftmax.forward +print(output) +``` diff --git a/docs/zh_cn/advanced_tutorials/registry.md b/docs/zh_cn/advanced_tutorials/registry.md index 47b7edcac4..89a8de854b 100644 --- a/docs/zh_cn/advanced_tutorials/registry.md +++ b/docs/zh_cn/advanced_tutorials/registry.md @@ -74,7 +74,7 @@ print(ACTIVATION.module_dict) ``` ```{note} -只有模块所在的文件被导入时,注册机制才会被触发,所以我们需要在某处导入该文件或者使用 `custom_imports` 字段动态导入该模块进而触发注册机制,详情见[导入自定义 Python 模块](config.md)。 +只有模块所在的文件被导入时,注册机制才会被触发,所以我们需要在某处导入该文件或者使用 `custom_imports` 字段动态导入该模块进而触发注册机制,详情见[导入自定义 Python 模块](config.md#导入自定义-python-模块)。 ``` 模块成功注册后,我们可以通过配置文件使用这个激活模块。