【mmengine】注册器(register) （进阶）继承mmengine注册器来实现自己项目的注册器

一、 进阶用法

为了方便跨库调用，MMEngine 提供了 22 个根注册器：

RUNNERS: Runner 的注册器
RUNNER_CONSTRUCTORS: Runner 的构造器
LOOPS: 管理训练、验证以及测试流程，如 EpochBasedTrainLoop
HOOKS: 钩子，如 CheckpointHook, ParamSchedulerHook
DATASETS: 数据集
DATA_SAMPLERS: DataLoader 的 Sampler，用于采样数据
TRANSFORMS: 各种数据预处理，如 Resize, Reshape
MODELS: 模型的各种模块
MODEL_WRAPPERS: 模型的包装器，如 MMDistributedDataParallel，用于对分布式数据并行
WEIGHT_INITIALIZERS: 权重初始化的工具
OPTIMIZERS: 注册了 PyTorch 中所有的 Optimizer 以及自定义的 Optimizer
OPTIM_WRAPPER: 对 Optimizer 相关操作的封装，如 OptimWrapper，AmpOptimWrapper
OPTIM_WRAPPER_CONSTRUCTORS: optimizer wrapper 的构造器
PARAM_SCHEDULERS: 各种参数调度器，如 MultiStepLR
METRICS: 用于计算模型精度的评估指标，如 Accuracy
EVALUATOR: 用于计算模型精度的一个或多个评估指标
TASK_UTILS: 任务强相关的一些组件，如 AnchorGenerator, BboxCoder
VISUALIZERS: 管理绘制模块，如 DetVisualizer 可在图片上绘制预测框
VISBACKENDS: 存储训练日志的后端，如 LocalVisBackend, TensorboardVisBackend
LOG_PROCESSORS: 控制日志的统计窗口和统计方法，默认使用 LogProcessor，如有特殊需求可自定义 LogProcessor
FUNCTIONS: 注册了各种函数，如 Dataloader 中传入的 collate_fn
INFERENCERS: 注册了各种任务的推理器，如 DetInferencer，负责检测任务的推理

二、 调用父节点的模块

MMEngine 中定义模块 RReLU，并往 MODELS 根注册器注册

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('调用 RReLU.forward')
        return x

假设有个行车项目叫 Driving_SHUAI，它也定义了 MODELS，并设置其父节点为 MMEngine 的 MODELS，这样就建立了层级结构。

from mmengine import Registry, MODELS as MMENGINE_MODELS

MODELS_SHUAI = Registry('model_SHUAI', parent=MMENGINE_MODELS, scope='Driving_SHUAI', locations=['Driving_SHUAI.models'])

下图是 MMEngine 和 Driving_SHUAI的注册器层级结构。

可以调用 count_registered_modules 函数打印已注册到 MMEngine 的模块以及层级结构。

from mmengine.registry import count_registered_modules

count_registered_modules()

在 Driving_SHUAI中定义模块 LogSoftmax_SHUAI，并往 Driving_SHUAI的 MODELS_SHUAI注册。

@MODELS_SHUAI.register_module()
class LogSoftmax_SHUAI(nn.Module):
    def __init__(self, dim=None):
        super().__init__()
        print('调用 LogSoftmax_SHUAI.__init__')

    def forward(self, x):
        print('调用 LogSoftmax_SHUAI.forward')
        return x

在 Driving_SHUAI 中使用字符串配置调用 LogSoftmax_SHUAI

model_softmax = MODELS_SHUAI.build(cfg=dict(type='LogSoftmax_SHUAI'))

也可以在 Driving_SHUAI中调用父节点 MMEngine 的模块。

model = MODELS.build(cfg=dict(type='RReLU', lower=0.2))
# 也可以加 scope
model = MODELS.build(cfg=dict(type='mmengine.RReLU'))

# 测试一些mmengine中的RReLU
import torch
input = torch.randn(2)
print(model(input))

如果不加前缀，build 方法首先查找当前节点是否存在该模块，如果存在则返回该模块，否则会继续向上查找父节点甚至祖先节点直到找到该模块，因此，如果当前节点和父节点存在同一模块并且希望调用父节点的模块，我们需要指定 scope 前缀。