MMDetection-自定义组件注册表(Registry)分析--以如何增加注意力模块为例

Troyella

已于 2022-07-14 14:12:52 修改

阅读量2.8k

点赞数 6

文章标签： python 机器学习深度学习目标检测 opencv

于 2022-07-01 14:35:14 首次发布

本文链接：https://blog.csdn.net/qq_43106412/article/details/125558140

版权

注册表是用来干什么的？

MMDection的所有功能都是基于注册器来实现的。注册器的使用目的是为了在算法训练、调参中通过直接更改配置文件(超参数，如学习率等)，完成对整个网络的修改。注册器在这期间可以帮助我们完成字符串到类的映射，自动解析config中的内容。其实说白了注册器只是通过一个类完成了string类型到类名的映射。
mmdetection的组件大多数以类的形式定义，下面是在mmdet\models\builder.py中实例化的注册器:

# Copyright (c) OpenMMLab. All rights reserved.
import warnings
from mmcv.cnn import MODELS as MMCV_MODELS
from mmcv.utils import Registry
MODELS = Registry('models', parent=MMCV_MODELS)

BACKBONES = MODELS
NECKS = MODELS
ROI_EXTRACTORS = MODELS
SHARED_HEADS = MODELS
HEADS = MODELS
LOSSES = MODELS
DETECTORS = MODELS

BACKBONES： 即模型的主干网络,用以对图片进行特征抽取.如常用的ResNeXt,HRNet等。
NECKS： 即对主干网络产生的特征图做一些特定的处理,如fpn多尺度特征融合信息。
Heads： 即目标检测的头部,包含了目标检测的主要算法逻辑,包括bbox的产生,回归target的计算,loss的计算等。
LOSS： 即损失函数的定义
DETECTOR： 将前面所有组件搭建而成的一个整体,通过加载detector来运行整体算法
PIPELINES： 数据增强管道类.定义了数据预处理和后处理部分

如何使用注册表？

在MMDetection中，注册一个模块的操作通过python装饰器方法实现。如果我们自定义了一个backbone，我们将这样使用Registry类的register_module装饰函数，也就是只需要一行代码，即可将自定义backbone注册到BACKBONES注册表中：

@BACKBONES.register_module
class Customize(nn.Module):

当我们完成一个自定义backbone的注册后，在配置文件做出超参数修改(以ResNet为例)：

backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        style='pytorch')

当我们在训练过程中，框架把上述的这个dict和BACKBONES注册表类传入build_from_cfg()函数，得到一个实例化的组件类，即通过 ‘type’ 为ResNet找到resnet类，初始化超参数depth、num_stages、out_indices。

通过实例解析自定义模块步骤(以创建SE_resnet为例，Neck等模块同理)

**步骤1：**在mmdetection/mmdet/models/backbone文件下创建一个SE_resnet.py文件，并在文件中添加@BACKBONES.register_module语句注册当前模块。

from __future__ import print_function, division, absolute_import
from collections import OrderedDict
from ..registry import BACKBONES
from mmcv.runner import load_checkpoint
import logging
from mmcv.cnn import constant_init, kaiming_init
import math
 
import torch.nn as nn
from torch.utils import model_zoo
 
__all__ = ['SENet', 'senet154', 'se_resnet50', 'se_resnet101', 'se_resnet152',
           'se_resnext50_32x4d', 'se_resnext101_32x4d']
 
pretrained_settings = {
            
    'senet154': {
            
        'imagenet': {
            
            'url': 'http://data.lip6.fr/cadene/pretrainedmodels/senet154-c7b49a05.pth',
            'input_space': 'RGB',
            'input_size': [3, 224, 224],
            'input_range': [0, 1],
            'mean': [0.485, 0.456, 0.406],
            'std': [0.229, 0.224, 0.225],
            'num_classes': 1000
        }
    },
    'se_resnet50': {
            
        'imagenet': {
            
            'url': 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet50-ce0d4300.pth',
            'input_space': 'RGB',
            'input_size': [3, 224, 224],
            'input_range': [0, 1],
            'mean': [0.485, 0.456, 0.406],
            'std': [0.229, 0.224, 0.225],
            'num_classes': 1000
        }
    },
    'se_resnet101': {
            
        'imagenet': {
            
            'url': 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet101-7e38fcc6.pth',
            'input_space': 'RGB',
            'input_size': [3, 224, 224],
            'input_range': [0, 1],
            'mean': [0.485, 0.456, 0.406],
            'std': [0.229, 0.224, 0.225],
            'num_classes': 1000
        }
    },
    'se_resnet152': {
            
        'imagenet': {
            
            'url': 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet152-d17c99b7.pth',
            'input_space': 'RGB',
            'input_size': [3, 224, 224],
            'input_range': [0, 1],
            'mean': [0.485, 0.456, 0.406],
            'std': [0.229, 0.224, 0.225],
            'num_classes': 1000
        }
    },
    'se_resnext50_32x4d': {
            
        'imagenet': {
            
            'url': 'http://data.lip6.fr/cadene/pretrainedmodels/se_