Background
霖的博客

MMCV入门：注册器和执行器

学习MMCV深度学习Python
2165 字
8 分钟阅读

本文基于 mmcv-full==1.6.0

配置

Config 类用于操作配置文件，它支持从多种文件格式中加载配置，包括 python, json 和 yaml。

目前支持以下四个预定义变量：

  • {{ fileDirname }} - 当前打开文件的目录名，例如 /home/your-username/your-project/folder
  • {{ fileBasename }} - 当前打开文件的文件名，例如 file.ext
  • {{ fileBasenameNoExtension }} - 当前打开文件不包含扩展名的文件名，例如 file
  • {{ fileExtname }} - 当前打开文件的扩展名，例如 .ext

不含重复键值对从基类配置文件继承

a = 1
b = dict(b1=[0, 1, 2], b2=None)
_base_ = 'config_a.py'
c = dict(b1=[3, 4, 5])

含重复键值对从基类配置文件继承

_base_ = 'config_a.py'
b = dict(b1=[3, 4, 5])

config_c.py 中的 b1 键值对会覆盖 config_a.py 中的 b1 键值对，而 b2 键值对会被保留。

从多个基类配置文件继承（基类配置文件不应包含相同的键）

c = (1, 2)
d = 'string'
_base_ = ['config_a.py', 'config_d.py']

从具有忽略字段的配置文件继承

_base_ = 'config_a.py'
b = dict(_delete_=True, b2=None, b3='new')

config_f.py 中的 _delete_ 字段会删除 config_a.py 中的 b 字典，然后添加新的键值对 b2 和 b3。

从基类引用变量

_base_ = 'config_a.py'
item = dict(item1={{ _base_.a }}, {{  _base_.b.b1 }})

config_g.py 中的 item 字典会引用 config_a.py 中的 a 和 b.b1 的值。

注册器

MMCV 使用注册器来管理具有相似功能的不同模块, 例如, 检测器中的主干网络、头部、和模型颈部。注册器可以看作类或函数到字符串的映射。 一个注册器中的类或函数通常有相似的接口，但是可以实现不同的算法或支持不同的数据集。 借助注册器，用户可以通过使用相应的字符串查找类或函数，并根据他们的需要实例化对应模块或调用函数获取结果。

使用 registry（注册器）管理代码库中的模型，需要以下三个步骤：

  1. 创建一个构建方法（可选，在大多数情况下您可以只使用默认方法）
  2. 创建注册器
  3. 使用此注册器来管理模块

Registry（注册器）的参数 build_func（构建函数） 用来自定义如何实例化类的实例或如何调用函数获取结果，默认使用build_from_cfg。

例

实现一个简单的注册器：

from mmcv.utils import Registry
# 创建转换器（converter）的注册器（registry）
CONVERTERS = Registry('converter')

然后可以在包中使用不同的转换器

from .builder import CONVERTERS


# 使用注册器管理模块
@CONVERTERS.register_module()
class Converter1(object):
    def __init__(self, a, b):
        self.a = a
        self.b = b
from .builder import CONVERTERS
from .converter1 import Converter1


# 使用注册器管理模块
@CONVERTERS.register_module()
def converter2(a, b)
    return Converter1(a, b)

通过这种方式就能建立字符串和其对应类或函数的映射。 之后就能通过配置文件使用转换器了。

converter1_cfg = dict(type='Converter1', a=a_value, b=b_value)
converter2_cfg = dict(type='converter2', a=a_value, b=b_value)
converter1 = CONVERTERS.build(converter1_cfg)
# returns the calling result
result = CONVERTERS.build(converter2_cfg)

自定义构建函数

from mmcv.utils import Registry


# 创建一个构建函数
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


# 创建一个用于转换器（converters）的注册器，并传递（registry）``build_converter`` 函数
CONVERTERS = Registry('converter', build_func=build_converter)

该功能类似于默认的build_from_cfg。在大多数情况下，默认就足够了。

注册器的层结构

OpenMMLab 下游代码库中所有 MODELS 注册器都是MMCV MODELS 注册器的子注册器。 可以使用以下两种方法从子注册器或相邻兄弟注册器构建模块：

例如，在 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
  1. 从子注册器中构建 我们可以通过以下代码在 MMDetection 或 MMClassification 中构建两个网络： 
    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. 从父注册器中构建 
    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'))

执行器

  • 支持以 EpochBasedRunner 和 IterBasedRunner 为单位的迭代模式以满足不同场景
  • 支持定制工作流以满足训练过程中各状态自由切换，目前支持训练和验证两个工作流。工作流可以简单理解为一个完成的训练和验证迭代过程。
  • 配合各类默认和自定义 Hook，对外提供了灵活扩展能力

EpochBasedRunner

EpochBasedRunner 是指以 epoch 为周期的工作流，例如设置 workflow = [(‘train’, 2), (‘val’, 1)] 表示循环迭代地训练 2 个 epoch，然后验证 1 个 epoch。MMDetection 目标检测框架默认采用的是 EpochBasedRunner。

其抽象逻辑如下所示：

# 训练终止条件
while curr_epoch < max_epochs:
    # 遍历用户设置的工作流，例如 workflow = [('train', 2)，('val', 1)]
    for i, flow in enumerate(workflow):
        # mode 是工作流函数，例如 train, epochs 是迭代次数
        mode, epochs = flow
        # 要么调用 self.train()，要么调用 self.val()
        epoch_runner = getattr(self, mode)
        # 运行对应工作流函数
        for _ in range(epochs):
            epoch_runner(data_loaders[i], **kwargs)

目前支持训练和验证两个工作流，以训练函数为例，其抽象逻辑是：

# epoch_runner 目前可以是 train 或者 val
def train(self, data_loader, **kwargs):
    # 遍历 dataset，共返回一个 epoch 的 batch 数据
    for i, data_batch in enumerate(data_loader):
        self.call_hook('before_train_iter')
        # 验证时候 train_mode=False
        self.run_iter(data_batch, train_mode=True, **kwargs)
        self.call_hook('after_train_iter')
   self.call_hook('after_train_epoch')

IterBasedRunner

IterBasedRunner 是指以 iter 为周期的工作流，例如设置 workflow = [(‘train’, 2)， (‘val’, 1)] 表示循环迭代的训练 2 个 iter，然后验证 1 个 iter，MMSegmentation 语义分割框架默认采用的是 IterBasedRunner。

其抽象逻辑如下所示：

# 虽然是 iter 单位，但是某些场合需要 epoch 信息，由 IterLoader 提供
iter_loaders = [IterLoader(x) for x in data_loaders]
# 训练终止条件
while curr_iter < max_iters:
    # 遍历用户设置的工作流，例如 workflow = [('train', 2)， ('val', 1)]
    for i, flow in enumerate(workflow):
        # mode 是工作流函数，例如 train, iters 是迭代次数
        mode, iters = flow
        # 要么调用 self.train()，要么调用 self.val()
        iter_runner = getattr(self, mode)
        # 运行对应工作流函数
        for _ in range(iters):
            iter_runner(iter_loaders[i], **kwargs)

目前支持训练和验证两个工作流，以验证函数为例，其抽象逻辑是：

# iter_runner 目前可以是 train 或者 val
def val(self, data_loader, **kwargs):
    # 获取 batch 数据，用于一次迭代
    data_batch = next(data_loader)
    self.call_hook('before_val_iter')
    outputs = self.model.val_step(data_batch, self.optimizer, **kwargs)
    self.outputs = outputs
    self.call_hook('after_val_iter')

开启训练的流程

  1. dataloader、model 和优化器等类初始化 
    # 模型类初始化
    model=...
    # 优化器类初始化
    # 典型值 cfg.optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001)
    optimizer = build_optimizer(model, cfg.optimizer)
    # 工作流对应的 dataloader 初始化
    data_loaders = [
            build_dataloader(
                ds,
                cfg.data.samples_per_gpu,
                cfg.data.workers_per_gpu,
                ...) for ds in dataset
        ]
  2. runner 类初始化 
    runner = build_runner(
        # cfg.runner 典型配置为
        # runner = dict(type='EpochBasedRunner', max_epochs=200)
        cfg.runner,
        default_args=dict(
            model=model,
            batch_processor=None,
            optimizer=optimizer,
            logger=logger))
  3. 注册默认训练所必须的 hook，和用户自定义 hook 
    # 注册定制必需的 hook
    runner.register_training_hooks(
        # lr相关配置，典型为
        # lr_config = dict(policy='step', step=[100, 150])
        cfg.lr_config,
        # 优化相关配置，例如 grad_clip 等
        optimizer_config,
        # 权重保存相关配置，典型为
        # checkpoint_config = dict(interval=1)，每个单位都保存权重
        cfg.checkpoint_config,
        # 日志相关配置
        cfg.log_config,
        ...)
    

    # 注册用户自定义 hook
    # 例如想使用 ema 功能，则可以设置 custom_hooks=[dict(type='EMAHook')]
    if cfg.get('custom_hooks', None):
        custom_hooks = cfg.custom_hooks
        for hook_cfg in cfg.custom_hooks:
            hook_cfg = hook_cfg.copy()
            priority = hook_cfg.pop('priority', 'NORMAL')
            hook = build_from_cfg(hook_cfg, HOOKS)
            runner.register_hook(hook, priority=priority)
  4. 开启训练 
    # workflow 典型为 workflow = [('train', 1)]
    # 此时就真正开启了训练
    runner.run(data_loaders, cfg.workflow)
MMCV入门：注册器和执行器
/blog/26037593
作者
霖
发布于
2026/3/18
许可协议
CC BY-NC-SA 4.0