Example #1
0
    def forward(self, features, im_info, boxes=None):
        assert len(features) == config.num_anchor_scales

        # get the predict results
        pred_cls_score_list = []
        pred_bbox_offsets_list = []
        all_anchors_list = []
        for x in features:
            t = F.relu(self.rpn_conv(x))
            pred_cls_score_list.append(self.rpn_cls_score(t))
            pred_bbox_offsets_list.append(self.rpn_bbox_offsets(t))
        # get anchors
        max_fm_stride = 2**(len(features) + 1)
        fm_stride = max_fm_stride
        anchor_scales = config.anchor_scales[::-1]
        for i in range(config.num_anchor_scales):
            anchor_scale = [anchor_scales[i]]
            layer_anchors = generate_anchors_opr(features[i], fm_stride,
                                                 anchor_scale,
                                                 config.anchor_aspect_ratios,
                                                 config.anchor_base_size)
            fm_stride = fm_stride // 2
            all_anchors_list.append(layer_anchors)
        # sample from the predictions
        with torch.no_grad():
            rpn_rois, rpn_rois_inds = find_top_rpn_proposals(
                self.training, pred_bbox_offsets_list, pred_cls_score_list,
                all_anchors_list, im_info)

        if self.training:
            with torch.no_grad():
                rpn_labels, rpn_bbox_targets = fpn_anchor_target(
                    boxes, im_info, all_anchors_list)
            pred_cls_score, pred_bbox_offsets = fpn_rpn_reshape(
                pred_cls_score_list, pred_bbox_offsets_list)

            # rpn loss
            valid_masks = rpn_labels >= 0
            objectness_loss = F.binary_cross_entropy_with_logits(
                pred_cls_score[valid_masks],
                rpn_labels[valid_masks].to(torch.float32),
                reduction="sum")
            pos_masks = rpn_labels == 1
            localization_loss = smooth_l1_loss(pred_bbox_offsets[pos_masks],
                                               rpn_bbox_targets[pos_masks],
                                               config.rpn_smooth_l1_beta,
                                               reduction="sum")
            normalizer = 1.0 / (config.train_batch_per_gpu *
                                config.num_sample_anchors)
            loss_rpn_cls = objectness_loss * normalizer  # cls: classification loss
            loss_rpn_loc = localization_loss * normalizer
            loss_dict = {}
            loss_dict['loss_rpn_cls'] = loss_rpn_cls
            loss_dict['loss_rpn_loc'] = loss_rpn_loc
            return rpn_rois, rpn_rois_inds, loss_dict
        else:
            return rpn_rois, rpn_rois_inds
Example #2
0
    def forward(self, features, im_info, boxes=None):
        # prediction
        pred_cls_score_list = []
        pred_bbox_offsets_list = []
        for x in features:
            t = F.relu(self.rpn_conv(x))
            pred_cls_score_list.append(self.rpn_cls_score(t))
            pred_bbox_offsets_list.append(self.rpn_bbox_offsets(t))
        # get anchors
        all_anchors_list = []
        fm_stride = 2**(len(features) + 1)
        for fm in features:
            layer_anchors = self.anchors_generator(fm, fm_stride)
            fm_stride = fm_stride // 2
            all_anchors_list.append(layer_anchors)
        # sample from the predictions
        rpn_rois, rpn_probs = find_top_rpn_proposals(self.training,
                                                     pred_bbox_offsets_list,
                                                     pred_cls_score_list,
                                                     all_anchors_list, im_info)

        if self.training:
            rpn_labels, rpn_bbox_targets = fpn_anchor_target(
                boxes, im_info, all_anchors_list)
            #rpn_labels = rpn_labels.astype(np.int32)
            pred_cls_score, pred_bbox_offsets = fpn_rpn_reshape(
                pred_cls_score_list, pred_bbox_offsets_list)

            # rpn loss
            valid_masks = rpn_labels >= 0
            valid_inds = mask_to_inds(valid_masks)
            objectness_loss = softmax_loss(pred_cls_score.ai[valid_inds],
                                           rpn_labels.ai[valid_inds])
            #objectness_loss = objectness_loss * valid_masks

            pos_masks = rpn_labels > 0
            localization_loss = smooth_l1_loss(pred_bbox_offsets,
                                               rpn_bbox_targets,
                                               config.rpn_smooth_l1_beta)
            localization_loss = localization_loss * pos_masks
            normalizer = 1.0 / (valid_masks.sum())
            loss_rpn_cls = objectness_loss.sum() * normalizer
            loss_rpn_loc = localization_loss.sum() * normalizer
            loss_dict = {}
            loss_dict['loss_rpn_cls'] = loss_rpn_cls
            loss_dict['loss_rpn_loc'] = loss_rpn_loc
            return rpn_rois, loss_dict
        else:
            return rpn_rois