mmdetection(pytorch0.4.1版本)模型构建部分源码解析(转)

我不是女神ヾ 2023-07-09 06:23 4阅读 0赞

在mmdetection中，实现了许多现有two-stage目标检测方法以及one-stage目标检测方法，且包含完整的数据载入、模型构建、模型训练以及模型测试部分的源码。因此非常适合在此基础上扩展实现其他目标检测算法。关于数据载入、模型训练以及模型测试部分的源码，在上一篇博客中有详细介绍。后期打算在mmdetection基础上添加其他算法。因此了解模型构建这部分源码十分重要。在mmdetecion官方[github][]介绍模型主要由四部分组成：backbone, neck, head, 以及ROI extractor。  
  
阅读mmdetection源码最简单地主线就是根据模型配置文件configs/\*.py文件中model部分去阅读。在mmdetection中，two-stage类目标检测方法model部分内容通常包含：backbone，neck，rpn\_head，bbox\_roi\_extractor和bbox\_head共五个部分。当然，在maskrcnn中还包括mask\_roi\_extractor和mask\_head。以maskrcnn模型的配置文件中model部分为例，介绍模型构建部分，下面是maskrcnn模型中model部分配置：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
</pre> </td> 
   <td> <pre>model = dict(
    type='MaskRCNN',  #模型名，对应模型声明在mmdet/modelsp/detectors/mask_rcnn.py文件中
    pretrained='modelzoo://resnet50', #backbone的预训练模型
    backbone=dict(
        type='ResNet',
        depth=50, 
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1, 
        style='pytorch', 
        dcn=dict(        
            modulated=False,
            deformable_groups=1,
            fallback_on_stride=False),
        stage_with_dcn=(False, True, True, True)),
    neck=dict(                   
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        num_outs=5),
    rpn_head=dict(
        type='RPNHead',
        in_channels=256,
        feat_channels=256,
        anchor_scales=[8],
        anchor_ratios=[0.5, 1.0, 2.0],
        anchor_strides=[4, 8, 16, 32, 64],
        target_means=[.0, .0, .0, .0],
        target_stds=[1.0, 1.0, 1.0, 1.0],
        use_sigmoid_cls=True),
    bbox_roi_extractor=dict(
        type='SingleRoIExtractor',
        roi_layer=dict(type='RoIAlign', out_size=7, sample_num=2),
        out_channels=256,
        featmap_strides=[4, 8, 16, 32]),
    bbox_head=dict(
        type='SharedFCBBoxHead',
        num_fcs=2,
        in_channels=256,
        fc_out_channels=1024,
        roi_feat_size=7,
        num_classes=81,
        target_means=[0., 0., 0., 0.],
        target_stds=[0.1, 0.1, 0.2, 0.2],
        reg_class_agnostic=False),
    mask_roi_extractor=dict(
        type='SingleRoIExtractor',
        roi_layer=dict(type='RoIAlign', out_size=14, sample_num=2),
        out_channels=256,
        featmap_strides=[4, 8, 16, 32]),
    mask_head=dict(
        type='FCNMaskHead',
        num_convs=4,
        in_channels=256,
        conv_out_channels=256,
        num_classes=81))
</pre> </td> 
  </tr> 
 </tbody> 
</table>

**注意**：

>  *  上面的pretrained参数只是针对backbone部分的预训练权重，在mmdetection中除了支持pytorch的modelzoo之外，mmdetection也提供了部分预训练权重，具体可以查看mmcv/runner/checkpoint.py文件。
>  *  如果想对整个模型进行预训练，比如我使用coco数据集预训练模型去训练pascal voc，则可以在configs/\*.py文件中load\_from参数，来指定coco预训练模型权重文件路径。

### **Backbone && neck** ###

目前官方实现了backbone包括：resnet,resnext以及ssd\_vgg。作者在resnet基础上还实现了dcn。至于neck部分目前只有FPN，而且对于two-stage方法，neck只能设置为FPN，对应代码实现部分在mmdet/models/necks/fpn.py。backbone和neck相关的配置参数如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
</pre> </td> 
   <td> <pre>backbone=dict(
    type='ResNet',
    depth=50, 
    num_stages=4,
    out_indices=(0, 1, 2, 3),
    frozen_stages=1,  #由于stage1的feature map比较大，因此冻结住，训练阶段可以节省内存
    style='pytorch',  #预训练模型载入来源，可以是caffe
    dcn=dict(         #dcn参数配置，如果仅仅使用普通resnet50，可以设置dcn=None
        modulated=False,
        deformable_groups=1,
        fallback_on_stride=False),
    stage_with_dcn=(False, True, True, True))，
neck=dict(                   
    type='FPN',
    in_channels=[256, 512, 1024, 2048], #对应于resnet的4个stage（C2,C3,C4,C5）输出通道数
    out_channels=256,                   #对应(P2,P3,P4,P5,P6)的输出通道数
    num_outs=5),                        #5个输出，对应P2-P6
</pre> </td> 
  </tr> 
 </tbody> 
</table>

backbone部分代码实现部分在mmdet/models/backbones/\*.py文件中。neck部分代码实现部分在mmdet/models/necks/fpn.py文件中。这两部分代码的实现不难，不予介绍，FPN结构如下图所示。值得注意的是，RetinaNet模型中neck部分也是采用的FPN，对应neck部分配置却不一样，RetinaNet模型中neck部分参数如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
</pre> </td> 
   <td> <pre>neck=dict(
    type='FPN',
    in_channels=[256, 512, 1024, 2048],
    out_channels=256,
    start_level=1,
    add_extra_convs=True,
    num_outs=5)
</pre> </td> 
  </tr> 
 </tbody> 
</table>

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xpc2h1aXdhbmc_size_16_color_FFFFFF_t_70][]  
额外添加了start\_level和add\_extra\_convs两个参数。start\_level=1表示从C3开始，因为RetinaNet模型中FPN的输出并不是P2-P6这5个输出，在RetinaNet中FPN的输出采用的是P3-P7这5个输出。在maskrcnn的FPN结构中是直接通过对P5进行maxpool得到P6，而在RetinaNet中，P6通过对P5进行stride=2卷积操作得到的，而P7则是同样对P6进行stride=2卷积操作得到。这一部分的区别，可以看源码mmdet/models/necks/fpn.py文件中第72行。

### **rpn\_head** ###

这一部分主要涉及到anchor的相关操作，在mmdet/models/anchor\_heads文件夹下实现了多种anchor\_heads，其中包括基类AnchorHead(实现在anchor\_head.py文件中），RetinaNet模型相关的RetinaHead类(实现在retina\_head.py文件中），two-stage方法的RPNHead类（实现在rpn\_head.py文件中）以及SSD方法SSDHead（实现在ssd\_head.py文件中）。后面三个类都继承自AnchorHead基类。在maskrcnn中使用的anchor\_head是RPNHead类，对应配置参数如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
</pre> </td> 
   <td> <pre>rpn_head=dict(
    type='RPNHead', 
    in_channels=256, # FPN中每个level的输出通道数，通常都为256
    feat_channels=256, # RPN的feature map的通道数
    anchor_scales=[8], # anchor的尺度，在使用FPN结构时，每个level只有一种尺度
    anchor_ratios=[0.5, 1.0, 2.0], #anchor比例：1:2,1:1,2:1这三种
    anchor_strides=[4, 8, 16, 32, 64], #对应于FPN中每个level的featuremap上每个pixel的感受野大小；
    target_means=[.0, .0, .0, .0],
    target_stds=[1.0, 1.0, 1.0, 1.0],
    use_sigmoid_cls=True) #anchor前景背景分类采用sigmoid，在早期的fasterrcnn采用的是softmax,在mmdetection中没有实现softmax分类
</pre> </td> 
  </tr> 
 </tbody> 
</table>

对应RPNHead类的实现如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
</pre> </td> 
   <td> <pre>class (AnchorHead): #继承自AnchorHead类

def __init__(self, in_channels, **kwargs):
        super(RPNHead, self).__init__(2, in_channels, **kwargs)

def _init_layers(self):
        self.rpn_conv = nn.Conv2d(
            self.in_channels, self.feat_channels, 3, padding=1)
        self.rpn_cls = nn.Conv2d(self.feat_channels,
                                 self.num_anchors * self.cls_out_channels, 1)
                                 
        self.rpn_reg = nn.Conv2d(self.feat_channels, self.num_anchors * 4, 1)

def init_weights(self):
        normal_init(self.rpn_conv, std=0.01)
        normal_init(self.rpn_cls, std=0.01)
        normal_init(self.rpn_reg, std=0.01)

def forward_single(self, x):  #表示输入FPN的一个level，比如P6输入，P6：1*256*13*16
        x = self.rpn_conv(x)      #conv不改变featuremap大小，且feat_channels=256，故x:1*256*13*16
        x = F.relu(x, inplace=True)
        rpn_cls_score = self.rpn_cls(x)  #对应于anchor的前景背景分类 1*3*13*16
        rpn_bbox_pred = self.rpn_reg(x)  #对应于anchor的回归 1*12*13*16
        return rpn_cls_score, rpn_bbox_pred

def loss(self, cls_scores, bbox_preds, gt_bboxes, img_metas, cfg):

losses = super(RPNHead, self).loss(cls_scores, bbox_preds, gt_bboxes,
                                           None, img_metas, cfg)
        return dict(
            loss_rpn_cls=losses['loss_cls'], loss_rpn_reg=losses['loss_reg'])
</pre> </td> 
  </tr> 
 </tbody> 
</table>

RPNHead类中实现的函数只是对AnchorHead类中部分函数的重写。由于RPNHead继承自AnchorHead类，在实际训练和测试中，直接调用的是还是继承自AnchorHead类的forward的函数。在RPN中通过在13x16的featuremap上生成13x16x3个anchors（注意，这里使用的是FPN结构，在不同level上，每个点只生成3个anchor，因此生成的anchors数是13x16x3）。经过后处理选择出proposal进行后续的bbox\_roi\_extractor。其中后处理操作包括：

>  *  删除背景框（13x16x3）；
>  *  NMS处理去除重合度大的框中分数低的那个；
>  *  取分数最大的K（top-k）个框用于后续bbox\_roi\_extractor；

这部分比较复杂，还涉及到一些RPN损失函数，比如smoth\_l1\_loss，focal loss以及delta2box的转换等等操作，直接看源码比较容易理解。关于Anchors这部分，[推荐博客][Link 1]讲解的比较清楚，推荐看。

### **ROI extractor** ###

在maskrcnn模型中，这部分涉及到box的操作以及mask的操作：bbox\_roi\_extractor和mask\_roi\_extractor。相关的配置如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
</pre> </td> 
   <td> <pre>bbox_roi_extractor=dict(
    type='SingleRoIExtractor',
    roi_layer=dict(type='RoIAlign', out_size=7, sample_num=2), #roi输出大小为7*7
    out_channels=256, #roi输出通道数
    featmap_strides=[4, 8, 16, 32]),  #对应4个level的proposal
mask_roi_extractor=dict(
    type='SingleRoIExtractor',
    roi_layer=dict(type='RoIAlign', out_size=14, sample_num=2),
    out_channels=256,
    featmap_strides=[4, 8, 16, 32]),
</pre> </td> 
  </tr> 
 </tbody> 
</table>

在上面的配置参数可以看出，bbox\_roi\_extractor和mask\_roi\_extractor的roi\_layer层类型都是SingleRoIExtractor。这个层的实现在mmdet/models/roi\_extractors/single\_level.py文件中。在FPN+RPN分别输出4个不同level的proposal，因此或生成4个对应的roi\_layer，每个roi\_layer输出的特征通道数out\_channels=256。对于box，每个roi的大小为7x7；对于mask，每个roi的输出大小14x14。

经过ROI extractor操作之后就是bbox的回归和分类以及mask的生成分类。这部分的配置参数如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
</pre> </td> 
   <td> <pre>bbox_head=dict(
    type='SharedFCBBoxHead',
    num_fcs=2, #2个全连接层
    in_channels=256, #roi输出特征通道数
    fc_out_channels=1024, #2个全连接层的节点数
    roi_feat_size=7,
    num_classes=81, #81个类别输出（80目标+1背景）
    target_means=[0., 0., 0., 0.],
    target_stds=[0.1, 0.1, 0.2, 0.2],
    reg_class_agnostic=False),
mask_head=dict(
    type='FCNMaskHead',
    num_convs=4, #4个卷积层
    in_channels=256, #roi输出特征通道数
    conv_out_channels=256, #4个卷积层的通道数均为256
    num_classes=81))    #mask的类别是81，
</pre> </td> 
  </tr> 
 </tbody> 
</table>

这一部分的实现对应于下图，该图截图自maskrcnn原文，如下图所示

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xpc2h1aXdhbmc_size_16_color_FFFFFF_t_70 1][]  
  
值得注意的是，在上面mask branch中倒数第二层应该是28x28x256的上采样层,在上面的参数配置中没有体现，但在mmdetection的代码有实现，可以看mmdet/models/mask\_heads/fcn\_mask\_head.py中FCNMaskHead默认的上采样层类型是deconv，层的实现在第58行。

### **模型构建** ###

上面根据模型配置文件结构介绍了构建maskrcnn模型的4个子模块，如何将这四个子模块拼接构成一个完整的maskrcnn模型？在mmdetecion中实例化一个具体模型，需要在mmdet/models/detectors/文件夹下中实现相应的模型，比如maskrcnn模型的实现在该文件夹下mask\_rcnn.py文件中，实现如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
</pre> </td> 
   <td> <pre>from .two_stage import TwoStageDetector
from ..registry import DETECTORS

@DETECTORS.register_module
class MaskRCNN(TwoStageDetector):

def __init__(self,
                 backbone,
                 neck,
                 rpn_head,
                 bbox_roi_extractor,
                 bbox_head,
                 mask_roi_extractor,
                 mask_head,
                 train_cfg,
                 test_cfg,
                 pretrained=None):
        super(MaskRCNN, self).__init__(
            backbone=backbone,
            neck=neck,
            rpn_head=rpn_head,
            bbox_roi_extractor=bbox_roi_extractor,
            bbox_head=bbox_head,
            mask_roi_extractor=mask_roi_extractor,
            mask_head=mask_head,
            train_cfg=train_cfg,
            test_cfg=test_cfg,
            pretrained=pretrained)
</pre> </td> 
  </tr> 
 </tbody> 
</table>

从上面可以看出maskrcnn模型的构建继承自TwoStageDetector类。TwoStageDetector类是一个通用的two-stage目标检测方法的基类，具体实现如下：

<table> 
 <tbody> 
  <tr> 
   <td> <pre>1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
</pre> </td> 
   <td> <pre>import torch
import torch.nn as nn

from .base import BaseDetector
from .test_mixins import RPNTestMixin, BBoxTestMixin, MaskTestMixin
from .. import builder
from ..registry import DETECTORS
from mmdet.core import bbox2roi, bbox2result, build_assigner, build_sampler

@DETECTORS.register_module
class TwoStageDetector(BaseDetector, RPNTestMixin, BBoxTestMixin,
                       MaskTestMixin):
#继承了四个基类BaseDetector（主要是一些模型常见训练，单多尺度测试相关操作的虚函数），RPNTestMixin，BBoxTestMixin，MaskTestMixin主要是模型单尺度/多尺度测试

def __init__(self,
                 backbone,
                 neck=None,
                 rpn_head=None,
                 bbox_roi_extractor=None,
                 bbox_head=None,
                 mask_roi_extractor=None,
                 mask_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        super(TwoStageDetector, self).__init__()
        self.backbone = builder.build_backbone(backbone)  #根据配置文件构架backbone

if neck is not None:
            self.neck = builder.build_neck(neck)        #根据配置文件构架neck,这里主要是FPN
        else:
            raise NotImplementedError

if rpn_head is not None:
            self.rpn_head = builder.build_head(rpn_head) #根据配置文件构架rpn
        #根据配置文件构架bbox_roi_extractor,bbox_head
        if bbox_head is not None:
            self.bbox_roi_extractor = builder.build_roi_extractor(
                bbox_roi_extractor)
            self.bbox_head = builder.build_head(bbox_head)
        #根据配置文件构架mask_roi_extractor,mask_head
        if mask_head is not None:
            self.mask_roi_extractor = builder.build_roi_extractor(
                mask_roi_extractor)
            self.mask_head = builder.build_head(mask_head)

self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        #模型参数权重初始化
        self.init_weights(pretrained=pretrained)

@property
    def with_rpn(self):
        return hasattr(self, 'rpn_head') and self.rpn_head is not None

def init_weights(self, pretrained=None):
        super(TwoStageDetector, self).init_weights(pretrained)
        self.backbone.init_weights(pretrained=pretrained)
        if self.with_neck:
            if isinstance(self.neck, nn.Sequential):
                for m in self.neck:
                    m.init_weights()
            else:
                self.neck.init_weights()
        if self.with_rpn:
            self.rpn_head.init_weights()
        if self.with_bbox:
            self.bbox_roi_extractor.init_weights()
            self.bbox_head.init_weights()
        if self.with_mask:
            self.mask_roi_extractor.init_weights()
            self.mask_head.init_weights()
    #提取特征（包括经过bakcbone,neck）
    def extract_feat(self, img):
        x = self.backbone(img)
        if self.with_neck:
            x = self.neck(x)
        return x

def forward_train(self,
                      img,
                      img_meta,
                      gt_bboxes,
                      gt_bboxes_ignore,
                      gt_labels,
                      gt_masks=None,
                      proposals=None):
        x = self.extract_feat(img)
      
        losses = dict()

# RPN 前向计算和损失
        if self.with_rpn:
            rpn_outs = self.rpn_head(x)
            rpn_loss_inputs = rpn_outs + (gt_bboxes, img_meta,
                                          self.train_cfg.rpn)
            rpn_losses = self.rpn_head.loss(*rpn_loss_inputs)
            losses.update(rpn_losses)

proposal_inputs = rpn_outs + (img_meta, self.test_cfg.rpn)
            proposal_list = self.rpn_head.get_bboxes(*proposal_inputs)
        else:
            proposal_list = proposals

# assign gts and sample proposals
        if self.with_bbox or self.with_mask:
            bbox_assigner = build_assigner(self.train_cfg.rcnn.assigner)
            bbox_sampler = build_sampler(
                self.train_cfg.rcnn.sampler, context=self)
            num_imgs = img.size(0)
            sampling_results = []
            for i in range(num_imgs):
                assign_result = bbox_assigner.assign(
                    proposal_list[i], gt_bboxes[i], gt_bboxes_ignore[i],
                    gt_labels[i])
                sampling_result = bbox_sampler.sample(
                    assign_result,
                    proposal_list[i],
                    gt_bboxes[i],
                    gt_labels[i],
                    feats=[lvl_feat[i][None] for lvl_feat in x])
                sampling_results.append(sampling_result)

# bbox head forward and loss
        if self.with_bbox:
            rois = bbox2roi([res.bboxes for res in sampling_results])
            # TODO: a more flexible way to decide which feature maps to use
            bbox_feats = self.bbox_roi_extractor(
                x[:self.bbox_roi_extractor.num_inputs], rois)
            cls_score, bbox_pred = self.bbox_head(bbox_feats)

bbox_targets = self.bbox_head.get_target(
                sampling_results, gt_bboxes, gt_labels, self.train_cfg.rcnn)
            loss_bbox = self.bbox_head.loss(cls_score, bbox_pred,
                                            *bbox_targets)
            losses.update(loss_bbox)

# mask head forward and loss
        if self.with_mask:
            pos_rois = bbox2roi([res.pos_bboxes for res in sampling_results])
            mask_feats = self.mask_roi_extractor(
                x[:self.mask_roi_extractor.num_inputs], pos_rois)
            mask_pred = self.mask_head(mask_feats)

mask_targets = self.mask_head.get_target(
                sampling_results, gt_masks, self.train_cfg.rcnn)
            pos_labels = torch.cat(
                [res.pos_gt_labels for res in sampling_results])
            loss_mask = self.mask_head.loss(mask_pred, mask_targets,
                                            pos_labels)
            losses.update(loss_mask)

return losses

def simple_test(self, img, img_meta, proposals=None, rescale=False):
        """Test without augmentation."""
        
    def aug_test(self, imgs, img_metas, rescale=False):
        """Test with augmentations.</pre> </td> 
  </tr> 
 </tbody> 
</table>

[github]: https://github.com/open-mmlab/mmdetection
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xpc2h1aXdhbmc_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/20200228090851221.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xpc2h1aXdhbmc=,size_16,color_FFFFFF,t_70
[Link 1]: https://www.cnblogs.com/q735613050/p/10631619.html
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xpc2h1aXdhbmc_size_16_color_FFFFFF_t_70 1]: https://img-blog.csdnimg.cn/20200228090805438.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xpc2h1aXdhbmc=,size_16,color_FFFFFF,t_70