mmdetection入门介绍-模型解析

拼搏现实的明天。 2023-06-27 05:59 16阅读 0赞

# 五、模型解析 #

我们可以参考一下docs/technical\_details.md里面的参考

模型回顾  
按照典型的约定，我们使用Dataset和DataLoader来加载具有多个工作程序的数据。Dataset返回数据项的字典，对应于模型的forward方法的参数。  
由于对象检测中的数据大小可能不同（图像大小，gt bbox大小等），  
在mmcv中引入了一种新的DataContainer类型，以帮助收集和分发不同大小的数据。

在mmdetection中，模型组件基本上分为4种类型。

*  backbone：通常是FCN网络，用于提取功能图，例如ResNet。
 *  neck：骨干和头部之间的部分，例如FPN，ASPP。
 *  head：用于特定任务的部分，例如bbox预测和mask预测。
 *  roi extractor：用于从 feature maps中提取feature的部分，例如RoI Align。

官方使用上述组件编写了一些通用的检测管道，例如“SingleStageDetector”和“ TwoStageDetector”。其中我们通过这两个内容来阅读一下框架，其中SingleStageDetector和TwoStageDetector均位于mmdet/models/detectors目录下的single\_stage.py和two\_stage.py.

## 5.1 单阶段检测 ##

可以看一下单阶段检测的源码

@DETECTORS.register_module
    class SingleStageDetector(BaseDetector):
        """Base class for single-stage detectors.
    
        Single-stage detectors directly and densely predict bounding boxes on the
        output features of the backbone+neck.
        """
    
        def __init__(self,
                     backbone,
                     neck=None,
                     bbox_head=None,
                     train_cfg=None,
                     test_cfg=None,
                     pretrained=None):
            super(SingleStageDetector, self).__init__()
            self.backbone = builder.build_backbone(backbone)
            if neck is not None:
                self.neck = builder.build_neck(neck)
            self.bbox_head = builder.build_head(bbox_head)
            self.train_cfg = train_cfg
            self.test_cfg = test_cfg
            self.init_weights(pretrained=pretrained)
    
        def init_weights(self, pretrained=None):
            super(SingleStageDetector, 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()
            self.bbox_head.init_weights()
    
        def extract_feat(self, img):
            """Directly extract features from the backbone+neck
            """
            x = self.backbone(img)
            if self.with_neck:
                x = self.neck(x)
            return x
    
        def forward_dummy(self, img):
            """Used for computing network flops.
    
            See `mmedetection/tools/get_flops.py`
            """
            x = self.extract_feat(img)
            outs = self.bbox_head(x)
            return outs
    
        def forward_train(self,
                          img,
                          img_metas,
                          gt_bboxes,
                          gt_labels,
                          gt_bboxes_ignore=None):
            x = self.extract_feat(img)
            outs = self.bbox_head(x)
            loss_inputs = outs + (gt_bboxes, gt_labels, img_metas, self.train_cfg)
            losses = self.bbox_head.loss(
                *loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
            return losses
    
        def simple_test(self, img, img_meta, rescale=False):
            x = self.extract_feat(img)
            outs = self.bbox_head(x)
            bbox_inputs = outs + (img_meta, self.test_cfg, rescale)
            bbox_list = self.bbox_head.get_bboxes(*bbox_inputs)
            bbox_results = [
                bbox2result(det_bboxes, det_labels, self.bbox_head.num_classes)
                for det_bboxes, det_labels in bbox_list
            ]
            return bbox_results[0]
    
        def aug_test(self, imgs, img_metas, rescale=False):
            raise NotImplementedError

从第一句话来看singlestageDetector传入参数为BaseDetector，也就是继承基础类来实现，然后就是初始化参数

def __init__(self,
                     backbone,
                     neck=None,
                     bbox_head=None,
                     train_cfg=None,
                     test_cfg=None,
                     pretrained=None):
            super(SingleStageDetector, self).__init__()
            self.backbone = builder.build_backbone(backbone)
            if neck is not None:
                self.neck = builder.build_neck(neck)
            self.bbox_head = builder.build_head(bbox_head)
            self.train_cfg = train_cfg
            self.test_cfg = test_cfg
            self.init_weights(pretrained=pretrained)

后面就是初始化权重，提取特征和用于计算网络触发

def init_weights(self, pretrained=None):
            super(SingleStageDetector, 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()
            self.bbox_head.init_weights()
    
        def extract_feat(self, img):
            """Directly extract features from the backbone+neck
            """
            x = self.backbone(img)
            if self.with_neck:
                x = self.neck(x)
            return x
    
        def forward_dummy(self, img):
            """Used for computing network flops.
    
            See `mmedetection/tools/get_flops.py`
            """
            x = self.extract_feat(img)
            outs = self.bbox_head(x)
            return outs

然后就是forward\_train，返回的是训练的loss

def forward_train(self,
                          img,
                          img_metas,
                          gt_bboxes,
                          gt_labels,
                          gt_bboxes_ignore=None):
            x = self.extract_feat(img)
            outs = self.bbox_head(x)
            loss_inputs = outs + (gt_bboxes, gt_labels, img_metas, self.train_cfg)
            losses = self.bbox_head.loss(
                *loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
            return losses

simple\_test输出的是bbox预处的结果

def simple_test(self, img, img_meta, rescale=False):
            x = self.extract_feat(img)
            outs = self.bbox_head(x)
            bbox_inputs = outs + (img_meta, self.test_cfg, rescale)
            bbox_list = self.bbox_head.get_bboxes(*bbox_inputs)
            bbox_results = [
                bbox2result(det_bboxes, det_labels, self.bbox_head.num_classes)
                for det_bboxes, det_labels in bbox_list
            ]
            return bbox_results[0]

其中，从函数的结果可以看出来bbox\_results输出的结果包含

*  det\_bboxes 检测的边界框
 *  det\_labels 检测物体的类别

## 5.2 双阶段检测 ##

还是先来看一下源码

@DETECTORS.register_module
    class TwoStageDetector(BaseDetector, RPNTestMixin, BBoxTestMixin,
                           MaskTestMixin):
        """Base class for two-stage detectors.
    
        Two-stage detectors typically consisting of a region proposal network and a
        task-specific regression head.
        """
    
        def __init__(self,
                     backbone,
                     neck=None,
                     shared_head=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)
    
            if neck is not None:
                self.neck = builder.build_neck(neck)
            if shared_head is not None:
                self.shared_head = builder.build_shared_head(shared_head)
            if rpn_head is not None:
                # RPN搭建后，有几张尺寸的特征图，执行几次其forward
                self.rpn_head = builder.build_head(rpn_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)
            if mask_head is not None:
                if mask_roi_extractor is not None:
                    self.mask_roi_extractor = builder.build_roi_extractor(
                        mask_roi_extractor)
                    self.share_roi_extractor = False
                else:
                    self.share_roi_extractor = True
                    self.mask_roi_extractor = self.bbox_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)
            # 训练阶段只有backbone是预训练初始化
            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_shared_head:
                self.shared_head.init_weights(pretrained=pretrained)
            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_head.init_weights()
                if not self.share_roi_extractor:
                    self.mask_roi_extractor.init_weights()
        # 特征图的提取层：backbone的输出+FPN的输出
        def extract_feat(self, img):
            """Directly extract features from the backbone+neck
            """
            x = self.backbone(img)# 经过backbone的前向计算
            if self.with_neck:
                x = self.neck(x) # neck FPN前向计算，得到输出特征图的tuple
            return x
    
        def forward_dummy(self, img):
            """Used for computing network flops.
    
            See `mmedetection/tools/get_flops.py`
            """
            outs = ()
            # backbone
            x = self.extract_feat(img)
            # rpn
            if self.with_rpn:
                rpn_outs = self.rpn_head(x)
                outs = outs + (rpn_outs, )
            proposals = torch.randn(1000, 4).cuda()
            # bbox head
            rois = bbox2roi([proposals])
            if self.with_bbox:
                bbox_feats = self.bbox_roi_extractor(
                    x[:self.bbox_roi_extractor.num_inputs], rois)
                if self.with_shared_head:
                    bbox_feats = self.shared_head(bbox_feats)
                cls_score, bbox_pred = self.bbox_head(bbox_feats)
                outs = outs + (cls_score, bbox_pred)
            # mask head
            if self.with_mask:
                mask_rois = rois[:100]
                mask_feats = self.mask_roi_extractor(
                    x[:self.mask_roi_extractor.num_inputs], mask_rois)
                if self.with_shared_head:
                    mask_feats = self.shared_head(mask_feats)
                mask_pred = self.mask_head(mask_feats)
                outs = outs + (mask_pred, )
            return outs
    
        # 训练过程的前向传播计算
        def forward_train(self,
                          img,
                          img_meta,
                          gt_bboxes,
                          gt_labels,
                          gt_bboxes_ignore=None,
                          gt_masks=None,
                          proposals=None):
            """
            Args:
                img (Tensor): of shape (N, C, H, W) encoding input images.
                    Typically these should be mean centered and std scaled.
    
                img_meta (list[dict]): list of image info dict where each dict has:
                    'img_shape', 'scale_factor', 'flip', and my also contain
                    'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
                    For details on the values of these keys see
                    `mmdet/datasets/pipelines/formatting.py:Collect`.
    
                gt_bboxes (list[Tensor]): each item are the truth boxes for each
                    image in [tl_x, tl_y, br_x, br_y] format.
    
                gt_labels (list[Tensor]): class indices corresponding to each box
    
                gt_bboxes_ignore (None | list[Tensor]): specify which bounding
                    boxes can be ignored when computing the loss.
    
                gt_masks (None | Tensor) : true segmentation masks for each box
                    used if the architecture supports a segmentation task.
    
                proposals : override rpn proposals with custom proposals. Use when
                    `with_rpn` is False.
    
            Returns:
                dict[str, Tensor]: a dictionary of loss components
            """
            x = self.extract_feat(img)# 下面这一步计算了前向的backbone传播和FPN
    
            # 从RPN开始有loss了
            losses = dict()
    
            # RPN forward and loss
            if self.with_rpn:
                rpn_outs = self.rpn_head(x) # 计算RPN前向传播，隐式调用forward注意遵照MRO顺序搜索
                # tuple可以直接作加法，相当于元组合并
                rpn_loss_inputs = rpn_outs + (gt_bboxes, img_meta,
                                              self.train_cfg.rpn)
                rpn_losses = self.rpn_head.loss(
                    *rpn_loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)# 计算前向loss，注意rpn的loss调用方法
                losses.update(rpn_losses) # 字典的合并方法
    
    
                proposal_cfg = self.train_cfg.get('rpn_proposal',
                                                  self.test_cfg.rpn)
                proposal_inputs = rpn_outs + (img_meta, proposal_cfg)  # 将RPN输出的box和相关设置信息输入proposal
                proposal_list = self.rpn_head.get_bboxes(*proposal_inputs) # 获得回归proposal
            else:
                proposal_list = proposals  # 直接指定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)
                if gt_bboxes_ignore is None:
                    gt_bboxes_ignore = [None for _ in range(num_imgs)]
                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)
                if self.with_shared_head:
                    bbox_feats = self.shared_head(bbox_feats)
                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:
                if not self.share_roi_extractor:
                    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)
                    if self.with_shared_head:
                        mask_feats = self.shared_head(mask_feats)
                else:
                    pos_inds = []
                    device = bbox_feats.device
                    for res in sampling_results:
                        pos_inds.append(
                            torch.ones(
                                res.pos_bboxes.shape[0],
                                device=device,
                                dtype=torch.uint8))
                        pos_inds.append(
                            torch.zeros(
                                res.neg_bboxes.shape[0],
                                device=device,
                                dtype=torch.uint8))
                    pos_inds = torch.cat(pos_inds)
                    mask_feats = bbox_feats[pos_inds]
                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."""
            assert self.with_bbox, "Bbox head must be implemented."
    
            x = self.extract_feat(img)
    
            proposal_list = self.simple_test_rpn(
                x, img_meta, self.test_cfg.rpn) if proposals is None else proposals
    
            det_bboxes, det_labels = self.simple_test_bboxes(
                x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)
            bbox_results = bbox2result(det_bboxes, det_labels,
                                       self.bbox_head.num_classes)
    
            if not self.with_mask:
                return bbox_results
            else:
                segm_results = self.simple_test_mask(
                    x, img_meta, det_bboxes, det_labels, rescale=rescale)
                return bbox_results, segm_results
    
        def aug_test(self, imgs, img_metas, rescale=False):
            """Test with augmentations.
    
            If rescale is False, then returned bboxes and masks will fit the scale
            of imgs[0].
            """
            # recompute feats to save memory
            proposal_list = self.aug_test_rpn(
                self.extract_feats(imgs), img_metas, self.test_cfg.rpn)
            det_bboxes, det_labels = self.aug_test_bboxes(
                self.extract_feats(imgs), img_metas, proposal_list,
                self.test_cfg.rcnn)
    
            if rescale:
                _det_bboxes = det_bboxes
            else:
                _det_bboxes = det_bboxes.clone()
                _det_bboxes[:, :4] *= img_metas[0][0]['scale_factor']
            bbox_results = bbox2result(_det_bboxes, det_labels,
                                       self.bbox_head.num_classes)
    
            # det_bboxes always keep the original scale
            if self.with_mask:
                segm_results = self.aug_test_mask(
                    self.extract_feats(imgs), img_metas, det_bboxes, det_labels)
                return bbox_results, segm_results
            else:
                return bbox_results

同样，我们还是来理清一下思路  
双阶段的检测同样是继承了BaseDetector, RPNTestMixin, BBoxTestMixin,MaskTestMixin来实现，典型的可以参考一下faster R-cnn。

首先是RPN网络来提取候选框(proposals),然后对这些proposals进行assign和采样，assign就是把候选框和gt进行匹配，区分出正负样本，如果是正样本还需要具体到哪一类（相当于粗略二分的基础上正样本还得分清楚），然后对这些候选框进行采样，再进行ROI\_pooling(ROI\_align)，得到适合fc的feature，最后进行回归和分类。

我们也可以从forward\_train看见这个过程

def forward_train(self,
                          img,
                          img_meta,
                          gt_bboxes,
                          gt_labels,
                          gt_bboxes_ignore=None,
                          gt_masks=None,
                          proposals=None):
    
            x = self.extract_feat(img)
            losses = dict()
            # RPN forward and loss
            if self.with_rpn:
                ...
            # assign gts and sample proposals
            if self.with_bbox or self.with_mask:
               ...
            # bbox head forward and loss
            if self.with_bbox:
                rois = bbox2roi([res.bboxes for res in sampling_results])
               ...
            # mask head forward and loss
            if self.with_mask:
               ...

然后就是simple\_test，和单阶段的类似

def simple_test(self, img, img_meta, proposals=None, rescale=False):
            """Test without augmentation."""
            assert self.with_bbox, "Bbox head must be implemented."
    
            x = self.extract_feat(img)
    
            proposal_list = self.simple_test_rpn(
                x, img_meta, self.test_cfg.rpn) if proposals is None else proposals
    
    
            det_bboxes, det_labels = self.simple_test_bboxes(
                x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)
            bbox_results = bbox2result(det_bboxes, det_labels,
                                       self.bbox_head.num_classes)
    
            if not self.with_mask:
                return bbox_results
            else:
                segm_results = self.simple_test_mask(
                    x, img_meta, det_bboxes, det_labels, rescale=rescale)
                return bbox_results, segm_results

其中，输出结果为bbox\_results, segm\_results，同样从函数的结果可以看出来bbox\_results输出的结果包含

*  det\_bboxes 检测的边界框
 *  det\_labels 检测物体的类别

其他参考：  
这里是mmdetection入门介绍 [前言][Link 1] 部分  
这里是mmdetection入门介绍 [test.py解析][test.py] 部分  
这里是mmdetection入门介绍 [train.py解析][train.py] 部分  
这里是mmdetection入门介绍 [模型解析][Link 2] 部分

[Link 1]: https://blog.csdn.net/klaus_x/article/details/103831820
[test.py]: https://blog.csdn.net/klaus_x/article/details/103831885
[train.py]: https://blog.csdn.net/klaus_x/article/details/103831939
[Link 2]: https://blog.csdn.net/klaus_x/article/details/103831987