注册器机制Registry
在众多深度学习开源库的代码中经常出现Registry代码块，例如OpenMMlab，facebookresearch和BasicSR中都使用了注册器机制。这块的代码经常会让新使用这些库的初学者感到一头雾水，本篇博客来分析一下注册器机制的原理与好处。
1. 为什么使用registry
在讲解registry原理前，我们先介绍一下，为何使用registry。registry的中文翻译是注册器。对于一个好用的深度学习代码库来说，通常都会内置多种损失函数，多种网络结构，以及多种优化器等。同时这类的库一般都支持从配置文件中，直接解析出模型结构与训练策略。那么如何优雅的从配置文件解析到具体的代码实现呢？这就是引入注册操作的意义，简而言之，注册器是为了方便找到相关模块。
2. registry代码阅读
在实现上不同代码库略有差异，但原理相同，所以这里就以BasicSR为例。
class Registry():
"""
The registry that provides name -> object mapping, to support third-party
users' custom modules.
To create a registry (e.g. a backbone registry):
.. 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):
"""
Args:
name (str): the name of this registry
"""
self._name = name
self._obj_map = {}
def _do_register(self, name, obj, suffix=None):
if isinstance(suffix, str):
name = name + '_' + suffix
assert (name not in self._obj_map), (f"An object named '{name}' was already registered "
f"in '{self._name}' registry!")
self._obj_map[name] = obj
def register(self, obj=None, suffix=None):
"""
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):
name = func_or_class.__name__
self._do_register(name, func_or_class, suffix)
return func_or_class
return deco
# used as a function call
name = obj.__name__
self._do_register(name, obj, suffix)
def get(self, name, suffix='basicsr'):
ret = self._obj_map.get(name)
if ret is None:
ret = self._obj_map.get(name + '_' + suffix)
print(f'Name {name} is not found, use name: {name}_{suffix}!')
if ret is None:
raise KeyError(f"No object named '{name}' found in '{self._name}' registry!")
return ret
def __contains__(self, name):
return name in self._obj_map
def __iter__(self):
return iter(self._obj_map.items())
def keys(self):
return self._obj_map.keys()
DATASET_REGISTRY = Registry('dataset')
ARCH_REGISTRY = Registry('arch')
MODEL_REGISTRY = Registry('model')
LOSS_REGISTRY = Registry('loss')
METRIC_REGISTRY = Registry('metric')
上面的代码为数据集，架构，网络，损失以及度量方式都创建了一个注册器对象。核心代码在register函数里，register函数使用了装饰器的设计，也就是只要在功能模块前进行@xx.register()进行装饰，就会对原有功能模块进行注册，并且最终返回原始的功能模块，不修改其原有功能。
在更下层的_do_register()可以看到，这里使用的是一个字典来执行注册操作，记录的键值对分别是模块的名称以及模块本身。这样一来，读取配置文件中的模块字符串后，我们就能够直接通过函数名或者类名找到其具体实现。
@LOSS_REGISTRY.register()
class L1Loss(nn.Module):
"""L1 (mean absolute error, MAE) loss.
Args:
loss_weight (float): Loss weight for L1 loss. Default: 1.0.
reduction (str): Specifies the reduction to apply to the output.
Supported choices are 'none' | 'mean' | 'sum'. Default: 'mean'.
"""
def __init__(self, loss_weight=1.0, reduction='mean'):
super(L1Loss, self).__init__()
if reduction not in ['none', 'mean', 'sum']:
raise ValueError(f'Unsupported reduction mode: {reduction}. Supported ones are: {_reduction_modes}')
self.loss_weight = loss_weight
self.reduction = reduction
def forward(self, pred, target, weight=None, **kwargs):
"""
Args:
pred (Tensor): of shape (N, C, H, W). Predicted tensor.
target (Tensor): of shape (N, C, H, W). Ground truth tensor.
weight (Tensor, optional): of shape (N, C, H, W). Element-wise weights. Default: None.
"""
return self.loss_weight * l1_loss(pred, target, weight, reduction=self.reduction)
