1. MMDetection训练和测试流程解析
首先查看 faster_rcnn_r50_fpn.py 配置文件,可以发现有六个字段组成。注册的结构名称为 ‘FasterRCNN’ ,这个注册的结构可以在 mmdet/model/detectors/faster_rcnn.py 中找到,其是继承至 TwoStageDetector 父类。初始化的配置组成就是 faster_rcnn_r50_fpn.py 的配置组成。
faster_rcnn.py
@DETECTORS.register_module() class FasterRCNN(TwoStageDetector): """Implementation of `Faster R-CNN <https://arxiv.org/abs/1506.01497>`_""" def __init__(self, backbone, rpn_head, roi_head, train_cfg, test_cfg, neck=None, pretrained=None, init_cfg=None): super(FasterRCNN, self).__init__( backbone=backbone, neck=neck, rpn_head=rpn_head, roi_head=roi_head, train_cfg=train_cfg, test_cfg=test_cfg, pretrained=pretrained, init_cfg=init_cfg)
faster_rcnn_r50_fpn.py
# model settings(简化了详细信息) model = dict( type='FasterRCNN', backbone=dict( type='ResNet',...), neck=dict( type='FPN',...), rpn_head=dict( type='RPNHead',...), roi_head=dict( type='StandardRoIHead',...), train_cfg=dict( rpn=dict(...), rpn_proposal=dict(...), rcnn=dict(...)), test_cfg=dict( rpn=dict(...), rcnn=dict(...)) )
可以注意到,在Faster-rcnn的组成中,除了基本的backbone,neck,rpn_head,roi_head四个部分之外,还存在着训练和测试的配置文件train_cfg 与 test_cfg。对于backbone,neck,rpn_head,roi_head这四个部分,可以发现在配置文件中是有具体的注册模型使用,比如在这里的faster-cnn的backbone中是使用的ResNet,neck部分使用了FPN结构,另外的rpn_head,roi_head同样使用了一个指定的模块结构并进行了参数初始化进行使用,而这些指定的模块都可以在mmdet/models路径下的py文件中找到其具体的实现代码。
但是后续的train_cfg 与 test_cfg是没有指定具体的文件的,只有一些字段的组成:rpn,rpn_proposal,rcnn,首先就是探究这些字段是如何在函数中被调用的,先来直接查看父类TwoStageDetector。
在TwoStageDetector中的__init__函数中可以发现,为什么会在train_cfg中使用rpn,rpn_proposal,rcnn这些字段,是因为这是初始化的时候默认调用的。同时,会在训练的过程中使用,相关代码如下所示:
class TwoStageDetector(BaseDetector): def __init__(self, backbone, neck=None, rpn_head=None, roi_head=None, train_cfg=None, test_cfg=None, pretrained=None, init_cfg=None): super(TwoStageDetector, self).__init__(init_cfg) if pretrained: warnings.warn('DeprecationWarning: pretrained is deprecated, ' 'please use "init_cfg" instead') backbone.pretrained = pretrained self.backbone = build_backbone(backbone) if neck is not None: self.neck = build_neck(neck) # rpn_head会同时将test和train的rpn部分一同注册 if rpn_head is not None: rpn_train_cfg = train_cfg.rpn if train_cfg is not None else None rpn_head_ = rpn_head.copy() rpn_head_.update(train_cfg=rpn_train_cfg, test_cfg=test_cfg.rpn) # 使用部分 self.rpn_head = build_head(rpn_head_) # roi_head会同时将test和train的rcnn部分一同注册 if roi_head is not None: rcnn_train_cfg = train_cfg.rcnn if train_cfg is not None else None roi_head.update(train_cfg=rcnn_train_cfg) # 使用部分 roi_head.update(test_cfg=test_cfg.rcnn) # 使用部分 roi_head.pretrained = pretrained self.roi_head = build_head(roi_head) self.train_cfg = train_cfg self.test_cfg = test_cfg ... def forward_train(self,..., *kwargs): x = self.extract_feat(img) losses = dict() # RPN forward and loss if self.with_rpn: proposal_cfg = self.train_cfg.get('rpn_proposal', self.test_cfg.rpn) rpn_losses, proposal_list = self.rpn_head.forward_train( proposal_cfg=proposal_cfg, ...) losses.update(rpn_losses) else: proposal_list = proposals roi_losses = self.roi_head.forward_train(...) losses.update(roi_losses) return losses
再重新观察faster_rcnn_r50_fpn.py文件,发现无论是rpn_head还是roi_head,都需要在train_cfg中为其分配一个assigner和sampler。
这里根据官方资料简单介绍一下:
- BBox Assigner
正负样本属性分配模块作用是进行正负样本定义或者正负样本分配(可能也包括忽略样本定义),正样本就是常说的前景样本(可以是任何类别),负样本就是背景样本。因为目标检测是一个同时进行分类和回归的问题,对于分类场景必然需要确定正负样本,否则无法训练。该模块至关重要,不同的正负样本分配策略会带来显著的性能差异,目前大部分目标检测算法都会对这个部分进行改进,至关重要。对应的代码在mmdet/core/bbox/assigners中
- BBox Sampler
在确定每个样本的正负属性后,可能还需要进行样本平衡操作。本模块作用是对前面定义的正负样本不平衡进行采样,力争克服该问题。一般在目标检测中 gt bbox 都是非常少的,所以正负样本比是远远小于 1 的。而基于机器学习观点:在数据极度不平衡情况下进行分类会出现预测倾向于样本多的类别,出现过拟合,为了克服该问题,适当的正负样本采样策略是非常必要的。对应的代码在mmdet/core/bbox/samplers中。
- BBox Encoder / Decoder
为了更好的收敛和平衡多个 loss,具体解决办法非常多,而 bbox 编解码策略也算其中一个,bbox 编码阶段对应的是对正样本的 gt bbox 采用某种编码变换(反操作就是 bbox 解码),最简单的编码是对 gt bbox 除以图片宽高进行归一化以平衡分类和回归分支,对应的代码在mmdet/core/bbox/coder中。
训练时候进行了编码,那么对应的测试环节需要进行解码。根据编码的不同,解码也是不同的。举个简单例子:假设训练时候对宽高是直接除以图片宽高进行归一化的,那么解码过程也仅仅需要乘以图片宽高即可。其代码和 bbox encoder 放在一起,在mmdet/core/bbox/coder中。
参考 Faster R-CNN 配置文件,这三者是作用与在同一个结构上的。
rpn_head=dict( type='RPNHead', ... bbox_coder=dict( type='DeltaXYWHBBoxCoder', target_means=[.0, .0, .0, .0], target_stds=[1.0, 1.0, 1.0, 1.0]), ... train_cfg=dict( rpn=dict( assigner=dict( type='MaxIoUAssigner', pos_iou_thr=0.7, neg_iou_thr=0.3, min_pos_iou=0.3, match_low_quality=True, ignore_iof_thr=-1), sampler=dict( type='RandomSampler', num=256, pos_fraction=0.5, neg_pos_ub=-1, add_gt_as_proposals=False), ...
1.1 训练流程
重新回到TwoStageDetector的代码中,对于two-stage的训练流程主要是调用forward_train这个函数,而在这个函数中,模块核心是调用 self.rpn_head.forward_train 和 self.roi_head.forward_train 函数,输出 losses 和其他相关数据。
#============= mmdet/models/detectors/two_stage.py/TwoStageDetector ============ def forward_train(...): # 先进行 backbone+neck 的特征提取 x = self.extract_feat(img) losses = dict() # RPN forward and loss if self.with_rpn: # 训练 RPN proposal_cfg = self.train_cfg.get('rpn_proposal', self.test_cfg.rpn) # 主要是调用 rpn_head 内部的 forward_train 方法 rpn_losses, proposal_list = self.rpn_head.forward_train(x,...) losses.update(rpn_losses) else: proposal_list = proposals # 第二阶段,主要是调用 roi_head 内部的 forward_train 方法 roi_losses = self.roi_head.forward_train(x, ...) losses.update(roi_losses) return losses
对于 one-stage 而言,具体如下所示:
#============= mmdet/models/detectors/single_stage.py/SingleStageDetector ============ def forward_train(...): super(SingleStageDetector, self).forward_train(img, img_metas) # 先进行 backbone+neck 的特征提取 x = self.extract_feat(img) # 主要是调用 bbox_head 内部的 forward_train 方法 losses = self.bbox_head.forward_train(x, ...) return losses
这个比 two-stage head 模块简单,因为其只有第一个 stage,对应的函数是 self.bbox_head.forward_train
1.2 测试流程
对于测试流程,如果是单尺度测试则会调用 TwoStageDetector中的 simple_test 方法,如果是多尺度测试,则调用 aug_test 方法,如下所示:
# 单尺度测试(est without augmentation) def simple_test(self, img, img_metas, proposals=None, rescale=False): assert self.with_bbox, 'Bbox head must be implemented.' x = self.extract_feat(img) if proposals is None: proposal_list = self.rpn_head.simple_test_rpn(x, img_metas) else: proposal_list = proposals return self.roi_head.simple_test( x, proposal_list, img_metas, rescale=rescale) # 多尺度测试(Test with augmentations) def aug_test(self, imgs, img_metas, rescale=False): x = self.extract_feats(imgs) proposal_list = self.rpn_head.aug_test_rpn(x, img_metas) return self.roi_head.aug_test( x, proposal_list, img_metas, rescale=rescale)
但是,可以看出在测试阶段,主要是调用了 Head 模块自身的 simple_test 或 aug_test 方法。
2. RPN_Head解析
这里以faster_rcnn_r50_fpn.py的配置文件为例,在这里配置文件的rpn使用了RPNHead这个类,所以现在来查看下这个类。
RPNHead继承自AnchorHead,同时AnchorHead又继承自BaseDenseHead与BBoxTestMixin。
@HEADS.register_module() class RPNHead(AnchorHead): ... @HEADS.register_module() class AnchorHead(BaseDenseHead, BBoxTestMixin): ... class BaseDenseHead(BaseModule, metaclass=ABCMeta):
2.1 训练流程
由于不断的继承,我们最终会发现,对于RPNHead来说,其训练函数入口forward_train最后是在其父类的父类BaseDenseHead中,其实现是在 mmdet/models/dense_heads/base_dense_head.py/BaseDenseHead 中,如下所示:
class BaseDenseHead(BaseModule, metaclass=ABCMeta): ... def forward_train(self, x, img_metas, gt_bboxes, gt_labels=None, gt_bboxes_ignore=None, proposal_cfg=None, **kwargs): # 调用各个子类实现的 forward 方法 outs = self(x) if gt_labels is None: loss_inputs = outs + (gt_bboxes, img_metas) else: loss_inputs = outs + (gt_bboxes, gt_labels, img_metas) # 调用各个子类实现的 loss 计算方法 losses = self.loss(*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore) if proposal_cfg is None: return losses else: # two-stage 算法还需要返回 proposal proposal_list = self.get_bboxes(*outs, img_metas, cfg=proposal_cfg) return losses, proposal_list
这里面有两个关键函数:self.loss 与 self(x)。每个算法的 Head 子类一般不会重写上述方法,但是每个 Head 子类都会重写 forward 和 loss 方法,其中 forward 方法用于运行 head 网络部分输出分类回归分支的特征图,而 loss 方法接收 forward 输出,并且结合 label 计算 loss。
而这里查看RPNHead这个类,就可以发现其重写了loss这个函数,不过调用的也是其父类AnchorHead的loss函数。
BaseDenseHead
BaseDenseHead 基类过于简单,对于 anchor-based 和 anchor-free 算法又进一步进行了继承,得到 AnchorHead 或者 AnchorFreeHead 类。在目前的各类算法实现中,绝大部分子类都是继承自 AnchorHead 或者 AnchorFreeHead,其提供了一些相关的默认操作,如果直接继承 BaseDenseHead 则子类需要重写大部分算法逻辑。
AnchorHead
刚刚分析代码到,在BaseDenseHead的forward_train训练入口函数中,都会实现子类的 forward 和 loss 方法。下面就对BaseDenseHead的子类AnchorHead的forward 和 loss 函数进行分析。其实也就说,在BaseDenseHead的forward_train中,self(x)就是调用AnchorHead的forward函数,self.loss就是调用AnchorHead的loss函数。
先来查看AnchorHead的forward函数:
# BBoxTestMixin 是多尺度测试时候调用 class AnchorHead(BaseDenseHead, BBoxTestMixin): # feats 是 backbone+neck 输出的多个尺度图 def forward(self, feats): # 对每张特征图单独计算预测输出 return multi_apply(self.forward_single, feats) # head 模块分类回归分支输出 def forward_single(self, x): cls_score = self.conv_cls(x) bbox_pred = self.conv_reg(x) return cls_score, bbox_pred
forward 函数比较简单,就是对多尺度特征图中每个特征图分别计算分类和回归输出即可,主要复杂度在 loss 函数中,其运行流程图如下所示:
代码如下:
class AnchorHead(BaseDenseHead, BBoxTestMixin): @force_fp32(apply_to=('cls_scores', 'bbox_preds')) def loss(self, cls_scores, bbox_preds, gt_bboxes, gt_labels, img_metas, gt_bboxes_ignore=None): """Compute losses of the head. Args: cls_scores (list[Tensor]): Box scores for each scale level Has shape (N, num_anchors * num_classes, H, W) bbox_preds (list[Tensor]): Box energies / deltas for each scale level with shape (N, num_anchors * 4, H, W) gt_bboxes (list[Tensor]): Ground truth bboxes for each image with shape (num_gts, 4) in [tl_x, tl_y, br_x, br_y] format. gt_labels (list[Tensor]): class indices corresponding to each box img_metas (list[dict]): Meta information of each image, e.g., image size, scaling factor, etc. gt_bboxes_ignore (None | list[Tensor]): specify which bounding boxes can be ignored when computing the loss. Default: None Returns: dict[str, Tensor]: A dictionary of loss components. """ featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores] assert len(featmap_sizes) == self.prior_generator.num_levels device = cls_scores[0].device anchor_list, valid_flag_list = self.get_anchors( featmap_sizes, img_metas, device=device) label_channels = self.cls_out_channels if self.use_sigmoid_cls else 1 cls_reg_targets = self.get_targets( anchor_list, valid_flag_list, gt_bboxes, img_metas, gt_bboxes_ignore_list=gt_bboxes_ignore, gt_labels_list=gt_labels, label_channels=label_channels) if cls_reg_targets is None: return None (labels_list, label_weights_list, bbox_targets_list, bbox_weights_list, num_total_pos, num_total_neg) = cls_reg_targets num_total_samples = ( num_total_pos + num_total_neg if self.sampling else num_total_pos) # anchor number of multi levels num_level_anchors = [anchors.size(0) for anchors in anchor_list[0]] # concat all level anchors and flags to a single tensor concat_anchor_list = [] for i in range(len(anchor_list)): concat_anchor_list.append(torch.cat(anchor_list[i])) all_anchor_list = images_to_levels(concat_anchor_list, num_level_anchors) losses_cls, losses_bbox = multi_apply( self.loss_single, cls_scores, bbox_preds, all_anchor_list, labels_list, label_weights_list, bbox_targets_list, bbox_weights_list, num_total_samples=num_total_samples) return dict(loss_cls=losses_cls, loss_bbox=losses_bbox)
分析:
- 在 loss 函数中首先会调用 get_anchors 函数得到默认 anchor 列表。而 get_anchors 函数内部会先计算多尺度特征图上每个特征点位置的 anchor,然后再计算有效 anchor 标志(因为在组织 batch 时候有些图片会进行左上角 padding,这部分像素人为加的,不需要考虑 anchor)
- 然后基于 anchor、gt bbox 以及其他必备信息调用 get_targets 函数计算每个预测分支对应的 target。get_targets 函数内部会调用 multi_apply(_get_targets_single) 函数对每张图片单独计算 target,而 _get_targets_single 函数实现的功能比较多,包括:bbox assigner、bbox sampler 和 bbox encoder 三个关键环节
- 在得到 targets 后,调用 loss_single 函数计算每个输出尺度的 loss 值,最终返回各个分支的 loss
AnchorFreeHead
AnchorFreeHead 逻辑比 AnchorHead 简单很多,主要是因为 anchor-free 类算法比 anchor-based 算法更加灵活多变,而且少了复杂的 anchor 生成过程,其 forward 方法实现和 AnchorHead 完全相同,而 loss 方法没有实现,其子类必须实现。
@HEADS.register_module()
class AnchorFreeHead(BaseDenseHead, BBoxTestMixin):
def forward(self, feats): return multi_apply(self.forward_single, feats)[:2] def forward_single(self, x): cls_feat = x reg_feat = x for cls_layer in self.cls_convs: cls_feat = cls_layer(cls_feat) cls_score = self.conv_cls(cls_feat) for reg_layer in self.reg_convs: reg_feat = reg_layer(reg_feat) bbox_pred = self.conv_reg(reg_feat) return cls_score, bbox_pred, cls_feat, reg_feat @abstractmethod @force_fp32(apply_to=('cls_scores', 'bbox_preds')) def loss(self, cls_scores, bbox_preds, gt_bboxes, gt_labels, img_metas, gt_bboxes_ignore=None): # loss 方法没有实现,其子类必须实现 raise NotImplementedError
上述内容介绍了在rpn_head中,具体的forward_train执行的内容。就是执行子类的forward与loss函数,下面就查看测试的流程。
2.2 测试过程
前面说过在测试流程中,最终会调用 Head 模块的 simple_test 或 aug_test 方法分别进行单尺度和多尺度测试,涉及到具体代码层面,one-stage 和 two-stage 调用函数有区别。具体上说,是调用head的simple_test_rpn 或 aug_test_rpn方法。但是最终调用的依然是 Head 模块的 get_bboxes 方法。
class TwoStageDetector(BaseDetector):
# 单尺度测试(est without augmentation) def simple_test(self, img, img_metas, proposals=None, rescale=False): if proposals is None: proposal_list = self.rpn_head.simple_test_rpn(x, img_metas) ... return self.roi_head.simple_test( x, proposal_list, img_metas, rescale=rescale) # 多尺度测试(Test with augmentations) def aug_test(self, imgs, img_metas, rescale=False): x = self.extract_feats(imgs) proposal_list = self.rpn_head.aug_test_rpn(x, img_metas) return self.roi_head.aug_test( x, proposal_list, img_metas, rescale=rescale)
AnchorHead
AnchorHead不仅集成了BaseDenseHead,同时还集成了BBoxTestMixin,而且其子类RPNHead的self.rpn_head.aug_test_rpn调用就是来自于BBoxTestMixin。但是该方法内部最终也是调用了 BaseDenseHead 中的get_bboxes 方法。同时,get_bboxes 方法又调用了子函数_get_bboxes_single方法。
class BBoxTestMixin(object):
... def simple_test_rpn(self, x, img_metas): """Test without augmentation, only for ``RPNHead`` and its variants, e.g., ``GARPNHead``, etc. Args: x (tuple[Tensor]): Features from the upstream network, each is a 4D-tensor. img_metas (list[dict]): Meta info of each image. Returns: list[Tensor]: Proposals of each image, each item has shape (n, 5), where 5 represent (tl_x, tl_y, br_x, br_y, score). """ rpn_outs = self(x) proposal_list = self.get_bboxes(*rpn_outs, img_metas=img_metas) # 归根到底还是调用了self.get_bboxes方法 return proposal_list def aug_test_rpn(self, feats, img_metas): """Test with augmentation for only for ``RPNHead`` and its variants, e.g., ``GARPNHead``, etc. Args: feats (tuple[Tensor]): Features from the upstream network, each is a 4D-tensor. img_metas (list[dict]): Meta info of each image. Returns: list[Tensor]: Proposals of each image, each item has shape (n, 5), where 5 represent (tl_x, tl_y, br_x, br_y, score). """ samples_per_gpu = len(img_metas[0]) aug_proposals = [[] for _ in range(samples_per_gpu)] for x, img_meta in zip(feats, img_metas): proposal_list = self.simple_test_rpn(x, img_meta) for i, proposals in enumerate(proposal_list): aug_proposals[i].append(proposals) # reorganize the order of 'img_metas' to match the dimensions # of 'aug_proposals' aug_img_metas = [] for i in range(samples_per_gpu): aug_img_meta = [] for j in range(len(img_metas)): aug_img_meta.append(img_metas[j][i]) aug_img_metas.append(aug_img_meta) # after merging, proposals will be rescaled to the original image size merged_proposals = [ merge_aug_proposals(proposals, aug_img_meta, self.test_cfg) for proposals, aug_img_meta in zip(aug_proposals, aug_img_metas) ] return merged_proposals
ps:这里再补充多一句,AnchorHead同时继承了BaseDenseHead, BBoxTestMixin两个父类,而get_bboxes方法是在BaseDenseHead父类中。
get_bboxes函数的主要流程:
- 遍历每个特征尺度输出分支,利用 nms_pre 配置参数对该层预测结果按照 scores 值进行从大到小进行 topk 截取,保留 scores 最高的前 nms_pre 的预测结果
- 对保留的预测结果进行 bbox 解码还原操作
- 还原到最原始图片尺度
- 如果需要进行 nms,则对所有分支预测保留结果进行统一 nms 即可,否则直接属于多尺度预测结果
ps:流程的1,2两步是get_bboxes调用_get_bboxes_single函数实现的,而流程的3,4两步是_get_bboxes_single函数再调用_bbox_post_process实现的。部分重要代码如下:
def _get_bboxes_single(self,...) # 遍历每个特征尺度输出分支 for level_idx, (cls_score, bbox_pred, score_factor, priors) in \ enumerate(zip(cls_score_list, bbox_pred_list, score_factor_list, mlvl_priors)): ... # 步骤1 results = filter_scores_and_topk( scores, cfg.score_thr, nms_pre, dict(bbox_pred=bbox_pred, priors=priors)) ... # 步骤2 bboxes = self.bbox_coder.decode( priors, bbox_pred, max_shape=img_shape) return self._bbox_post_process(mlvl_scores, mlvl_labels, mlvl_bboxes,...) def _bbox_post_process(self,...): # 步骤3 if rescale: mlvl_bboxes /= mlvl_bboxes.new_tensor(scale_factor) ... # 步骤4 if with_nms: det_bboxes, keep_idxs = batched_nms(mlvl_bboxes, mlvl_scores, mlvl_labels, cfg.nms)
AnchorFreeHead
AnchorFreeHead同样继承自BaseDenseHead, BBoxTestMixin,和AnchorHead的父类是完全一样的。所以其get_bboxes的所在位置与AnchorHead一样,也是在BaseDenseHead中。
除了 RPN 算法的多尺度测试是在mmdet/models/dense_heads/rpn_test_mixin.py,其余 Head 多尺度测试都是在 mmdet/models/dense_heads/dense_test_mixins.py/BBoxTestMixin 中实现,其思路是对多尺度图片中每张图片单独运行 get_bboxes,然后还原到原图尺度,最后把多尺度图片预测结果合并进行统一 nms。
以上,结束了Two-Stage中rpn_head部分的讲解,接下来就是rpn_head的下一步——roi_head。
3. ROI_Head解析
这里以faster_rcnn_r50_fpn.py的配置文件为例,在这里配置文件的rpn使用了StandardRoIHead这个类,所以现在来查看下这个类。
StandardRoIHead继承了3个父类,分别是BaseRoIHead, BBoxTestMixin, MaskTestMixin
3.1 训练流程
在two-stage网络的训练过程forward_train中,在实现了self.rpn_head.forward_train之后,随即而来的就是self.roi_head.forward_train。也就是训练完rpn_head,再进行roi_head的训练操作。而在faster-rcnn中作为roi_head的StandardRoIHead,其继承的父类BaseRoIHead本是只是定义了forward_train这个函数,但是没有任何操作,所以需要StandardRoIHead子类对其进行复写。其核心代码如下:
@HEADS.register_module() class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin): def forward_train(self, x, img_metas, proposal_list, gt_bboxes, gt_labels, ...): if self.with_bbox or self.with_mask: num_imgs = len(img_metas) sampling_results = [] for i in range(num_imgs): # 对每张图片进行 bbox 正负样本属性分配 assign_result = self.bbox_assigner.assign( proposal_list[i], ...) # 然后进行正负样本采样 sampling_result = self.bbox_sampler.sample( assign_result, proposal_list[i], ...) sampling_results.append(sampling_result) losses = dict() if self.with_bbox: # bbox 分支 forward,返回 loss bbox_results = self._bbox_forward_train(...) losses.update(bbox_results['loss_bbox']) if self.with_mask: # mask 分支 forward,返回 loss return losses def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels, img_metas): rois = bbox2roi([res.bboxes for res in sampling_results]) # forward bbox_results = self._bbox_forward(x, rois) # 计算 target bbox_targets = self.bbox_head.get_targets(...) # 计算 loss loss_bbox = self.bbox_head.loss(...) return ... def _bbox_forward(self, x, rois): # roi 提取 bbox_feats = self.bbox_roi_extractor( x[:self.bbox_roi_extractor.num_inputs], rois) # bbox head 网络前向 cls_score, bbox_pred = self.bbox_head(bbox_feats) return ...
从上述逻辑可以看出,StandardRoIHead 中 forward_train 函数仅仅是对内部的 bbox_head 相关函数进行调用,例如 get_targets 和 loss,本身 StandardRoIHead 类不做具体算法逻辑计算。也就是和rpn_head的算法逻辑不同,这里的roi_head的forward_train操作中是不涉及get_targets 和 loss函数。
可以参考 Faster R-CNN 配置文件理解 StandardRoIHead 和 bbox_head 的关系:
roi_head=dict( type='StandardRoIHead', bbox_roi_extractor=dict( type='SingleRoIExtractor', roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0), out_channels=256, featmap_strides=[4, 8, 16, 32]), bbox_head=dict( type='Shared2FCBBoxHead', in_channels=256, 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=1.0), loss_bbox=dict(type='L1Loss', loss_weight=1.0))))
StandardRoIHead 类包装了 bbox_roi_extractor 和 bbox_head 的实例,前者用于 RoI 特征提取,后者才是真正计算分类和回归的逻辑。在 bbox_head 中除了网络模型有些变换外,loss计算过程是非常类似的,其 get_targets 和 loss 计算过程都是封装在基类 mmdet/models/roi_heads/bbox_heads/bbox_head.py 中。
class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin): def _bbox_forward(self, x, rois): bbox_feats = self.bbox_roi_extractor(...) cls_score, bbox_pred = self.bbox_head(...) # 调用bbox_head本身的forward函数 ...... return bbox_results def _bbox_forward_train(self, x, ...): bbox_results = self._bbox_forward(x, rois) ... bbox_targets = self.bbox_head.get_targets(...) loss_bbox = self.bbox_head.loss(...) # 调用bbox_head本身的loss函数
在这里,StandardRoIHead使用的是Shared2FCBBoxHead这个bbox_head,现在直接来看这个类的代码,了解其forward函数与loss函数的实现。其中,Shared2FCBBoxHead继承自ConvFCBBoxHead,而ConvFCBBoxHead又继承自BBoxHead。ConvFCBBoxHead复写了BBoxHead中的forward函数,而其余的loss函数,get_targets,get_bboxes均是BBoxHead父类实现的。(在StandardRoIHead这里的get_bboxes好像没有被使用上)
3.2 测试流程
测试流程是调用 Head 模块的 simple_test 和 aug_test 函数。单尺度测试 bbox 相关实现代码在 mmdet/models/roi_heads/test_mixins.py/BBoxTestMixin 的 simple_test_bboxes 函数中。而多尺度测试的bbox 相关实现代码在 mmdet/models/roi_heads/test_mixins.py/BBoxTestMixin 的 aug_test_bboxes 函数中。这两个函数是直接在StandardRoIHead中的simple_test和aug_test中直接调用的,相关代码如下所示:
# 父类 class BBoxTestMixin: def simple_test_bboxes(self, ...): rois = bbox2roi(proposals) # roi 提取+ forward,输出预测结果 bbox_results = self._bbox_forward(x, rois) cls_score = bbox_results['cls_score'] bbox_pred = bbox_results['bbox_pred'] det_bboxes = [] det_labels = [] for i in range(len(proposals)): # 对预测结果进行解码输出 bbox 和对应 label det_bbox, det_label = self.bbox_head.get_bboxes(...) det_bboxes.append(det_bbox) det_labels.append(det_label) return det_bboxes, det_labels def aug_test_bboxes(self, ...): aug_bboxes = [] aug_scores = [] for x, img_meta in zip(feats, img_metas): ... # TODO more flexible proposals = bbox_mapping(proposal_list[0][:, :4], img_shape, scale_factor, flip, flip_direction) rois = bbox2roi([proposals]) bbox_results = self._bbox_forward(x, rois) bboxes, scores = self.bbox_head.get_bboxes(...) det_bboxes, det_labels = multiclass_nms(...) return det_bboxes, det_labels # 子类继承 @HEADS.register_module() class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin): def simple_test(self,...): det_bboxes, det_labels = self.simple_test_bboxes(...) ... if self.with_mask: segm_results = self.simple_test_mask(...) def aug_test(self, x, proposal_list, img_metas, rescale=False): det_bboxes, det_labels = self.aug_test_bboxes(...) ... if self.with_mask: segm_results = self.aug_test_mask(...)
实际上依然是调用了 Head 模块内部的 get_bboxes 函数,处理逻辑和 dense_head 差不多( 解码+还原尺度+ nms)。
参考资料:
1. 轻松掌握 MMDetection 中 Head 流程
2. 轻松掌握 MMDetection 整体构建流程(一)
3. 轻松掌握 MMDetection 整体构建流程(二)