MMDetection2（三）：自定义模型

目录

前言

一、 添加一个新的backbone

1、定义一个新的backbone（以MobileNet为例）

2、导入新的module

3、在配置文件中使用backbone

二、添加一个新的necks

1、定义一个neck（例如：PAFPN）

2、导入新的module

3、更改配置文件

三、添加一个新的heads

四、添加一个新的loss

---

前言

本篇主要介绍如何使用MMDetection2的组件搭建自己的的模型。（主要根据原英文文档翻译，有兴趣建议去看看原文档）

我们将模型的组件简单的分为5种类型：

• backbone: 通常使用一个FCN（全卷积神经网络）来提取feature maps(特征图),例如：ResNet, MobileNet.
• neck: 这个组件是放在backbones和heads之间的，例如：FPN(特征金字塔网络), PAFPN.
• head: 这个组件是用于一些指定的任务，例如：bbox prediction and mask prediction.
• roi extractor: 这个组件是用来从feature maps中提取ROI features的。例如：RoI Align.
• loss: 这个组件是放在head组件中用来计算损失的。例如： FocalLoss, L1Loss以及 GHMLoss.

一、 添加一个新的backbone

1、定义一个新的backbone（以MobileNet为例）

首先建立一个新的文件:mmdet/models/backbones/mobilenet.py，在该文件中实现MobileNet的细节:

import torch.nn as nn
from ..builder import BACKBONES


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

    def __init__(self, arg1, arg2):
        # 细节跳过
        pass

    def forward(self, x):  # should return a tuple
        # 细节跳过
        pass

    def init_weights(self, pretrained=None):
        # 细节跳过 
        pass

2、导入新的module

你可以在mmdet/models/backbones/__init__.py文件中添加以下代码：

from .mobilenet import MobileNet

 或者，你如果想要避免更改源码，可以选择在配置文件中添加以下代码：

# 这种方式我还没有试过，又机会试一下（因为不太清楚它指的配置文件是指哪一个orz）
custom_imports = dict(
    imports=['mmdet.models.backbones.mobilenet'],
    allow_failed_imports=False)

3、在配置文件中使用backbone

在configs/_base_/models/中选择相应模型进行backbone的替换:

model = dict(
    ...
    backbone=dict(
        type='MobileNet',
        arg1=xxx,
        arg2=xxx),
    ...

 

二、添加一个新的necks

1、定义一个neck（例如：PAFPN）

创建一个新文件mmdet/models/necks/pafpn.py。

# 官方现在已经将这个模块实现了
from ..builder import NECKS

@NECKS.register
class PAFPN(nn.Module):

    def __init__(self,
                in_channels,
                out_channels,
                num_outs,
                start_level=0,
                end_level=-1,
                add_extra_convs=False):
        pass

    def forward(self, inputs):
        # implementation is ignored
        pass

2、导入新的module

你可以在mmdet/models/necks/__init__.py文件中添加以下代码：

from .pafpn import PAFPN

  或者，你如果想要避免更改源码，可以选择在配置文件中添加以下代码：

custom_imports = dict(
    imports=['mmdet.models.necks.pafpn'],   # 官网这个地方写的是imports=['mmdet.models.necks.mobilenet'],应该是不小心写错了
    allow_failed_imports=False)

3、更改配置文件

在configs/_base_/models/中选择相应模型进行neck的替换:

neck=dict(
    type='PAFPN',
    in_channels=[256, 512, 1024, 2048],
    out_channels=256,
    num_outs=5)

三、添加一个新的heads

我们将以Double Head R-CNN作为例子来讲解如何构造一个新的head组件。

首先需要在mmdet/models/roi_heads/bbox_heads/中添加一个新的bbox head文件double_bbox_head.py（不过现在官方已经实现了这个）

Double Head R-CNN实现了一个新的bbox head用于目标检测任务。要实现一个bbox head，我们需要下面这个新module的三个函数：__init__()、init_weights()、forward()。

from mmdet.models.builder import HEADS
from .bbox_head import BBoxHead

@HEADS.register_module()
class DoubleConvFCBBoxHead(BBoxHead):
    r"""Bbox head used in Double-Head R-CNN

                                      /-> cls
                  /-> shared convs ->
                                      \-> reg
    roi features
                                      /-> cls
                  \-> shared fc    ->
                                      \-> reg
    """  # noqa: W605
    
    # 个人觉得这个__init__函数比较难定义，因为__init__函数输入的参数不知道要具体定义哪些
    def __init__(self,
                 num_convs=0,
                 num_fcs=0,
                 conv_out_channels=1024,
                 fc_out_channels=1024,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 **kwargs):
        kwargs.setdefault('with_avg_pool', True)
        super(DoubleConvFCBBoxHead, self).__init__(**kwargs)

    def init_weights(self):
        # conv layers are already initialized by ConvModule

    def forward(self, x_cls, x_reg):

 如果有必要的话，我们还需要实现一个ROI Head,我们准备通过继承StandardRoIHead来实现新的DoubleHeadRoIHead。我们可以看到官方已经实现了StrandardRoIHead这个类，它包含了以下几个函数:

import torch
from mmdet.core import bbox2result, bbox2roi, build_assigner, build_sampler
from ..builder import HEADS, build_head, build_roi_extractor
from .base_roi_head import BaseRoIHead
from .test_mixins import BBoxTestMixin, MaskTestMixin


@HEADS.register_module()
class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
    """Simplest base roi head including one bbox head and one mask head.
    """

    def init_assigner_sampler(self):
    def init_bbox_head(self, bbox_roi_extractor, bbox_head):
    def init_mask_head(self, mask_roi_extractor, mask_head):
    def init_weights(self, pretrained):
    def forward_dummy(self, x, proposals):
    def forward_train(self,
                      x,
                      img_metas,
                      proposal_list,
                      gt_bboxes,
                      gt_labels,
                      gt_bboxes_ignore=None,
                      gt_masks=None):
    # 我们在子类中主要对这个方法进行重写
    def _bbox_forward(self, x, rois):
    def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels,
                            img_metas):
    def _mask_forward_train(self, x, sampling_results, bbox_feats, gt_masks,
                            img_metas):
    def _mask_forward(self, x, rois=None, pos_inds=None, bbox_feats=None):
    def simple_test(self,
                    x,
                    proposal_list,
                    img_metas,
                    proposals=None,
                    rescale=False):
        """Test without augmentation."""

 Double Head的更改主要是在bbox_forward函数中，并且Double Head会继承StandardRoIHead的其他的方法，在mmdet/models/roi_heads/路径下创建double_roi_head.py文件，在其中我们实现新的ROI Head的代码如下所示:

from ..builder import HEADS
from .standard_roi_head import StandardRoIHead


@HEADS.register_module()
# 继承父类StandardRoIHead
class DoubleHeadRoIHead(StandardRoIHead):
    """RoI head for Double Head RCNN

    https://arxiv.org/abs/1904.06493
    """

    def __init__(self, reg_roi_scale_factor, **kwargs):
        super(DoubleHeadRoIHead, self).__init__(**kwargs)
        self.reg_roi_scale_factor = reg_roi_scale_factor

    # 只对这个bbox_forward进行重写
    def _bbox_forward(self, x, rois):
        # bbox_roi_extractor是将RPN获取的proposal，通过ROI方法对应到特征图上然后获取相应的特征
        bbox_cls_feats = self.bbox_roi_extractor(
            x[:self.bbox_roi_extractor.num_inputs], rois)
        bbox_reg_feats = self.bbox_roi_extractor(
            x[:self.bbox_roi_extractor.num_inputs],
            rois,
            roi_scale_factor=self.reg_roi_scale_factor)
        # 这个地方我还不太明白shared_head是一个什么模块
        if self.with_shared_head:
            bbox_cls_feats = self.shared_head(bbox_cls_feats)
            bbox_reg_feats = self.shared_head(bbox_reg_feats)
        # bbox_head是根据ROI层输出的特征获取对应的classification 结果 和 bbox结果的
        # 这个地方的self.bbox_head应该是一个double_bbox_head对象，在model的配置文件定义的
        cls_score, bbox_pred = self.bbox_head(bbox_cls_feats, bbox_reg_feats)

        bbox_results = dict(
            cls_score=cls_score,
            bbox_pred=bbox_pred,
            bbox_feats=bbox_cls_feats)
        return bbox_results

接着，我们需要在mmdet/models/bbox_heads/__init__.py和mmdet/models/roi_heads/__init__.py中添加相应的module。（跟上面的添加neck的方式类似）。

当然也可以选择在配置文件中添加下面这行代码，效果也是一样的。

custom_imports=dict(
    imports=['mmdet.models.roi_heads.double_roi_head', 'mmdet.models.bbox_heads.double_bbox_head'])

 Double Head R-CNN的model配置文件如下所示：

# 继承父类
_base_ = '../faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'

model = dict(
    # 只对roi_head模块进行更改 
    roi_head=dict(
        # 使用新建的"DoubleHeadRoIHead"
        # 这个地方的参数需要和DoubleHeadRoIHead类的初始化函数的参数保持一致
        type='DoubleHeadRoIHead',
        reg_roi_scale_factor=1.3,
        bbox_head=dict(
            #因为参数的名称有所不同，所以设置_delete_ = True，用新的字段名代替旧的字段名
            _delete_=True,
            type='DoubleConvFCBBoxHead',
            #这个地方的参数名都需要和DoubleConvFCBBoxHeadl类的初始化函数的参数保持一致
            num_convs=4,
            num_fcs=2,
            in_channels=256,
            conv_out_channels=1024,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=80,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=False,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=2.0),
            loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=2.0))))

从MMDetection2.0开始，配置文件支持使用继承，这样以便于使用者可以专注于具体的修改内容。

Double Head R-CNN主要使用一个新的DoubleHeadRoIHead以及一个新的DoubleConvFCBBoxHead，它们的参数都必须和每一个module中的__init__参数保持一致。

四、添加一个新的loss

待续.......