diff --git a/ppdet/modeling/heads/cascade_head.py b/ppdet/modeling/heads/cascade_head.py index ba17b3265b79740968ccea22c2cb55764ae26c37..fab310ef11931ff3a8ce4eb6014428b4998b339b 100644 --- a/ppdet/modeling/heads/cascade_head.py +++ b/ppdet/modeling/heads/cascade_head.py @@ -160,7 +160,8 @@ class CascadeHead(BBoxHead): bbox_weight=[[10., 10., 5., 5.], [20.0, 20.0, 10.0, 10.0], [30.0, 30.0, 15.0, 15.0]], num_cascade_stages=3, - bbox_loss=None): + bbox_loss=None, + reg_class_agnostic=True): nn.Layer.__init__(self, ) self.head = head self.roi_extractor = roi_extractor @@ -173,6 +174,9 @@ class CascadeHead(BBoxHead): self.num_cascade_stages = num_cascade_stages self.bbox_loss = bbox_loss + self.reg_class_agnostic = reg_class_agnostic + num_bbox_delta = 4 if reg_class_agnostic else 4 * num_classes + self.bbox_score_list = [] self.bbox_delta_list = [] for i in range(num_cascade_stages): @@ -190,7 +194,7 @@ class CascadeHead(BBoxHead): delta_name, nn.Linear( in_channel, - 4, + num_bbox_delta, weight_attr=paddle.ParamAttr(initializer=Normal( mean=0.0, std=0.001)))) self.bbox_score_list.append(bbox_score) @@ -227,6 +231,13 @@ class CascadeHead(BBoxHead): bbox_feat = self.head(rois_feat, i) scores = self.bbox_score_list[i](bbox_feat) deltas = self.bbox_delta_list[i](bbox_feat) + + # TODO (lyuwenyu) Is it correct for only one class ? + if not self.reg_class_agnostic and i < self.num_cascade_stages - 1: + deltas = deltas.reshape([-1, self.num_classes, 4]) + labels = scores[:, :-1].argmax(axis=-1) + deltas = deltas[paddle.arange(deltas.shape[0]), labels] + head_out_list.append([scores, deltas, rois]) pred_bbox = self._get_pred_bbox(deltas, rois, self.bbox_weight[i]) diff --git a/ppdet/modeling/proposal_generator/rpn_head.py b/ppdet/modeling/proposal_generator/rpn_head.py index a046fd73e90037ebb566c53b6824f50d2aea8e90..8a431eeac208a052ed8de5dfb7278948cfbcf042 100644 --- a/ppdet/modeling/proposal_generator/rpn_head.py +++ b/ppdet/modeling/proposal_generator/rpn_head.py @@ -68,6 +68,7 @@ class RPNHead(nn.Layer): derived by from_config """ __shared__ = ['export_onnx'] + __inject__ = ['loss_rpn_bbox'] def __init__(self, anchor_generator=_get_class_default_kwargs(AnchorGenerator), @@ -76,7 +77,8 @@ class RPNHead(nn.Layer): 12000, 2000), test_proposal=_get_class_default_kwargs(ProposalGenerator), in_channel=1024, - export_onnx=False): + export_onnx=False, + loss_rpn_bbox=None): super(RPNHead, self).__init__() self.anchor_generator = anchor_generator self.rpn_target_assign = rpn_target_assign @@ -91,6 +93,7 @@ class RPNHead(nn.Layer): self.train_proposal = ProposalGenerator(**train_proposal) if isinstance(test_proposal, dict): self.test_proposal = ProposalGenerator(**test_proposal) + self.loss_rpn_bbox = loss_rpn_bbox num_anchors = self.anchor_generator.num_anchors self.rpn_feat = RPNFeat(in_channel, in_channel) @@ -298,7 +301,12 @@ class RPNHead(nn.Layer): loc_tgt = paddle.concat(loc_tgt) loc_tgt = paddle.gather(loc_tgt, pos_ind) loc_tgt.stop_gradient = True - loss_rpn_reg = paddle.abs(loc_pred - loc_tgt).sum() + + if self.loss_rpn_bbox is None: + loss_rpn_reg = paddle.abs(loc_pred - loc_tgt).sum() + else: + loss_rpn_reg = self.loss_rpn_bbox(loc_pred, loc_tgt).sum() + return { 'loss_rpn_cls': loss_rpn_cls / norm, 'loss_rpn_reg': loss_rpn_reg / norm