MMCV-Registry

目录

什么是注册表

实例

自定义构造函数

多层级注册表

---

MMCV通过操作注册表实现对例如backbone，head和necks等在检测器中拥有相似功能的模块的管理。OpenMMLab中的很多工程通过使用注册表来管理数据集和模型中的模块。

什么是注册表

在MMCV中，注册表可被当做类到字符串的一种映射。这些被包含在单个注册表中的类通常有类似的API，但是执行不同的算法或者支持不同的数据集。有了注册表，用户可以通过其对应的字符串来查询和实例化类，并且使用他们所想使用的实例化模块。典型的例子就是在大多数OpenMMLab工程中的config系统，它通过configs使用注册表创建hooks，runners，models和datasets。涉及的API可参考这里

可通过以下三个步骤使用Registry在源码中管理模块。

1.创建一个构造方法（可选，大部分情况使用默认方法）；

2.创建一个注册表；

3.使用注册表管理模块。

Registry中的build_func参数是用来自定义如何实例化类，默认值是build_fron_cfg，执行过程参考这里

实例

这里我们将展示在库中使用注册表管理模块的简单例子。你可以在OpenMMLab中找到更多实际的例子。

假设我们想要执行一系列数据转换操作，将不同格式的数据转换成期望的数据格式。我们创建一个名为converters的文件作为库。在这个库中，我们首先创建一个文件来执行构造

converters/builders.py

from mmcv.utils import Registry
# create a registry for converters
CONVERTERS = Registry('converters')

然后我们在这个库中执行不同的转换操作。例如，我们在converters/converter1.py中执行Converter1

from .builder import CONVERTERS

# use the registry to manage the module
@CONVERTERS.register_module()
class Converter1(object):
    def __init__(self, a, b):
        self.a = a
        self.b = b

使用注册表管理模块的关键步骤是：创建模块的时候，通过@CONVERTERS.register_module()将被执行的模块注册到注册表CONVERTERS中。通过这种方式，CONVERTERS将在字符串和类之间建立和维持一种映射，如下所示：

'Converter1' -> <class 'Converter1'>

注意：注册表机制只有在模块所在文件被导入时才会触发。因此你需要在某些地方导入该文件。详情参考这里

如果该模块被成功注册，你可以通过如下方法，通过configs使用这个转换：

converter_cfg = dict(type='Converter1', a=a_value, b=b_value)
converter = CONVERTERS.build(converter_cfg)

自定义构造函数

假设我们想要自定义converters如何被构造，我们可以执行自定义build_func并且将它传入注册表。

from mmcv.utils import Registry

# create a build function
def build_converter(cfg, registry, *args, **kwargs):
    cfg_ = cfg.copy()
    converter_type = cfg_.pop('type')
    if converter_type not in registry:
        raise KeyError(f'Unrecognized converter type {converter_type}')
    else:
        converter_cls = registry.get(converter_type)

    converter = converter_cls(*args, **kwargs, **cfg_)
    return converter

# create a registry for converters and pass ``build_converter`` function
CONVERTERS = Registry('converter', build_func=build_converter)

注意：在该例子中，我们演示了如何使用build_func参数来自定义创建类实例的方法。这个功能类似于默认的build_from_cig。在多数情况下，默认值即可满足使用需求。build_model_from_cfg在nn.Sequential中创建Pytorch模块时也被执行，你可以直接使用它们，而不是自己执行。

多层级注册表

你可以从多个OpenMMLab框架中构造模块。例如，你可以使用MMClassification中所有的backbone在MMdetection中进行目标检测，也可以将目标检测模块和语义分割模块组合。

源码中所有下游的MODELS注册表都是MMCV的MODEL注册表的子注册表。通常有两种方法从子注册表或父注册表构造模块。

1.从子注册表构造

例如，在MMdetection中我们定义

from mmcv.utils import Registry
from mmcv.cnn import MODELS as MMCV_MODELS
MODELS = Registry('model', parent=MMCV_MODELS)

@MODELS.register_module()
class NetA(nn.Module):
    def forward(self, x):
        return x

在MMClassification中我们定义

from mmcv.utils import Registry
from mmcv.cnn import MODELS as MMCV_MODELS
MODELS = Registry('model', parent=MMCV_MODELS)

@MODELS.register_module()
class NetB(nn.Module):
    def forward(self, x):
        return x + 1

我们可以在MMD或MMC任意一个中构建两个网络：

from mmdet.models import MODELS
net_a = MODELS.build(cfg=dict(type='NetA'))
net_b = MODELS.build(cfg=dict(type='mmcls.NetB'))

或者

from mmcls.models import MODELS
net_a = MODELS.build(cfg=dict(type='mmdet.NetA'))
net_b = MODELS.build(cfg=dict(type='NetB'))

2.从父注册表构造

在MMCV中共享的MODEL注册表是所有下游注册表的父注册表（根注册表）

from mmcv.cnn import MODELS as MMCV_MODELS
net_a = MMCV_MODELS.build(cfg=dict(type='mmdet.NetA'))
net_b = MMCV_MODELS.build(cfg=dict(type='mmcls.NetB'))