def forward(self, fpn_fms, rcnn_rois, im_info=None, gt_boxes=None): rcnn_rois, labels, bbox_targets = self.get_ground_truth( rcnn_rois, im_info, gt_boxes) fpn_fms = [fpn_fms[x] for x in self.in_features] pool_features = layers.roi_pool( fpn_fms, rcnn_rois, self.stride, self.pooling_size, self.pooling_method, ) flatten_feature = F.flatten(pool_features, start_axis=1) roi_feature = F.relu(self.fc1(flatten_feature)) roi_feature = F.relu(self.fc2(roi_feature)) pred_logits = self.pred_cls(roi_feature) pred_offsets = self.pred_delta(roi_feature) if self.training: # loss for rcnn classification loss_rcnn_cls = F.loss.cross_entropy(pred_logits, labels, axis=1) # loss for rcnn regression pred_offsets = pred_offsets.reshape(-1, self.cfg.num_classes, 4) num_samples = labels.shape[0] fg_mask = labels > 0 loss_rcnn_bbox = layers.smooth_l1_loss( pred_offsets[fg_mask, labels[fg_mask] - 1], bbox_targets[fg_mask], self.cfg.rcnn_smooth_l1_beta, ).sum() / F.maximum(num_samples, mge.tensor(1)) loss_dict = { "loss_rcnn_cls": loss_rcnn_cls, "loss_rcnn_bbox": loss_rcnn_bbox, } return loss_dict else: # slice 1 for removing background pred_scores = F.softmax(pred_logits, axis=1)[:, 1:] pred_offsets = pred_offsets.reshape(-1, 4) target_shape = (rcnn_rois.shape[0], self.cfg.num_classes, 4) # rois (N, 4) -> (N, 1, 4) -> (N, 80, 4) -> (N * 80, 4) base_rois = F.broadcast_to( F.expand_dims(rcnn_rois[:, 1:5], axis=1), target_shape).reshape(-1, 4) pred_bbox = self.box_coder.decode(base_rois, pred_offsets) return pred_bbox, pred_scores
def forward(self, image, im_info, gt_boxes=None): image = self.preprocess_image(image) features = self.backbone(image) features = [features[f] for f in self.in_features] box_logits, box_offsets = self.head(features) box_logits_list = [ _.transpose(0, 2, 3, 1).reshape(image.shape[0], -1, self.cfg.num_classes) for _ in box_logits ] box_offsets_list = [ _.transpose(0, 2, 3, 1).reshape(image.shape[0], -1, 4) for _ in box_offsets ] anchors_list = self.anchor_generator(features) all_level_box_logits = F.concat(box_logits_list, axis=1) all_level_box_offsets = F.concat(box_offsets_list, axis=1) all_level_anchors = F.concat(anchors_list, axis=0) if self.training: gt_labels, gt_offsets = self.get_ground_truth( all_level_anchors, gt_boxes, im_info[:, 4].astype(np.int32), ) all_level_box_logits = all_level_box_logits.reshape( -1, self.cfg.num_classes) all_level_box_offsets = all_level_box_offsets.reshape(-1, 4) gt_labels = gt_labels.flatten() gt_offsets = gt_offsets.reshape(-1, 4) valid_mask = gt_labels >= 0 fg_mask = gt_labels > 0 num_fg = fg_mask.sum() gt_targets = F.zeros_like(all_level_box_logits) gt_targets[fg_mask, gt_labels[fg_mask] - 1] = 1 loss_cls = layers.sigmoid_focal_loss( all_level_box_logits[valid_mask], gt_targets[valid_mask], alpha=self.cfg.focal_loss_alpha, gamma=self.cfg.focal_loss_gamma, ).sum() / F.maximum(num_fg, 1) loss_bbox = layers.smooth_l1_loss( all_level_box_offsets[fg_mask], gt_offsets[fg_mask], beta=self.cfg.smooth_l1_beta, ).sum() / F.maximum(num_fg, 1) * self.cfg.loss_bbox_weight total = loss_cls + loss_bbox loss_dict = { "total_loss": total, "loss_cls": loss_cls, "loss_bbox": loss_bbox, } self.cfg.losses_keys = list(loss_dict.keys()) return loss_dict else: # currently not support multi-batch testing assert image.shape[0] == 1 pred_boxes = self.box_coder.decode(all_level_anchors, all_level_box_offsets[0]) pred_boxes = pred_boxes.reshape(-1, 4) scale_w = im_info[0, 1] / im_info[0, 3] scale_h = im_info[0, 0] / im_info[0, 2] pred_boxes = pred_boxes / F.concat( [scale_w, scale_h, scale_w, scale_h], axis=0) clipped_boxes = layers.get_clipped_boxes(pred_boxes, im_info[0, 2:4]).reshape( -1, 4) pred_score = F.sigmoid(all_level_box_logits)[0] return pred_score, clipped_boxes
def get_losses(self, anchors, pred_logits, pred_offsets, gt_boxes, im_info): # pylint: disable=too-many-statements def positive_bag_loss(logits, axis=1): weight = 1.0 / (1.0 - logits) weight /= weight.sum(axis=axis, keepdims=True) bag_prob = (weight * logits).sum(axis=1) return -layers.safelog(bag_prob) def negative_bag_loss(logits, gamma): return (logits**gamma) * (-layers.safelog(1.0 - logits)) pred_scores = F.sigmoid(pred_logits) box_prob_list = [] positive_losses = [] clamp_eps = 1e-7 bucket_size = self.cfg.bucket_size for bid in range(im_info.shape[0]): boxes_info = gt_boxes[bid, :im_info[bid, 4].astype("int32")] # id 0 is used for background classes, so -1 first labels = boxes_info[:, 4].astype("int32") - 1 pred_box = self.box_coder.decode(anchors, pred_offsets[bid]).detach() overlaps = layers.get_iou(boxes_info[:, :4], pred_box).detach() thresh1 = self.cfg.box_iou_threshold thresh2 = F.clip(overlaps.max(axis=1, keepdims=True), lower=thresh1 + clamp_eps, upper=1.0) gt_pred_prob = F.clip((overlaps - thresh1) / (thresh2 - thresh1), lower=0, upper=1.0) image_boxes_prob = F.zeros(pred_logits.shape[1:]).detach() # guarantee that nonzero_idx is not empty if gt_pred_prob.max() > clamp_eps: _, nonzero_idx = F.cond_take(gt_pred_prob != 0, gt_pred_prob) # since nonzeros is only 1 dim, use num_anchor to get real indices num_anchors = gt_pred_prob.shape[1] anchors_idx = nonzero_idx % num_anchors gt_idx = nonzero_idx // num_anchors image_boxes_prob[anchors_idx, labels[gt_idx]] = gt_pred_prob[gt_idx, anchors_idx] box_prob_list.append(image_boxes_prob) # construct bags for objects match_quality_matrix = layers.get_iou(boxes_info[:, :4], anchors).detach() num_gt = match_quality_matrix.shape[0] _, matched_idx = F.topk( match_quality_matrix, k=bucket_size, descending=True, no_sort=True, ) matched_idx = matched_idx.detach() matched_idx_flatten = matched_idx.reshape(-1) gather_idx = labels.reshape(-1, 1) gather_idx = F.broadcast_to(gather_idx, (num_gt, bucket_size)) gather_src = pred_scores[bid, matched_idx_flatten] gather_src = gather_src.reshape(num_gt, bucket_size, -1) matched_score = F.indexing_one_hot(gather_src, gather_idx, axis=2) topk_anchors = anchors[matched_idx_flatten] boxes_broad_cast = F.broadcast_to( F.expand_dims(boxes_info[:, :4], axis=1), (num_gt, bucket_size, 4)).reshape(-1, 4) matched_offsets = self.box_coder.encode(topk_anchors, boxes_broad_cast) reg_loss = layers.smooth_l1_loss( pred_offsets[bid, matched_idx_flatten], matched_offsets, beta=self.cfg.smooth_l1_beta).sum( axis=-1) * self.cfg.reg_loss_weight matched_reg_scores = F.exp(-reg_loss) positive_losses.append( positive_bag_loss(matched_score * matched_reg_scores.reshape(-1, bucket_size), axis=1)) num_foreground = im_info[:, 4].sum() pos_loss = F.concat(positive_losses).sum() / F.maximum( 1.0, num_foreground) box_probs = F.stack(box_prob_list, axis=0) neg_loss = negative_bag_loss( pred_scores * (1 - box_probs), self.cfg.focal_loss_gamma).sum() / F.maximum( 1.0, num_foreground * bucket_size) alpha = self.cfg.focal_loss_alpha pos_loss = pos_loss * alpha neg_loss = neg_loss * (1 - alpha) loss_dict = { "total_loss": pos_loss + neg_loss, "pos_loss": pos_loss, "neg_loss": neg_loss, } return loss_dict
def forward(self, features, im_info, boxes=None): # prediction features = [features[x] for x in self.in_features] # get anchors anchors_list = self.anchor_generator(features) pred_cls_logit_list = [] pred_bbox_offset_list = [] for x in features: t = F.relu(self.rpn_conv(x)) scores = self.rpn_cls_score(t) pred_cls_logit_list.append( scores.reshape( scores.shape[0], self.num_cell_anchors, scores.shape[2], scores.shape[3], )) bbox_offsets = self.rpn_bbox_offsets(t) pred_bbox_offset_list.append( bbox_offsets.reshape( bbox_offsets.shape[0], self.num_cell_anchors, 4, bbox_offsets.shape[2], bbox_offsets.shape[3], )) # get rois from the predictions rpn_rois = self.find_top_rpn_proposals(pred_cls_logit_list, pred_bbox_offset_list, anchors_list, im_info) if self.training: rpn_labels, rpn_offsets = self.get_ground_truth( anchors_list, boxes, im_info[:, 4].astype(np.int32)) pred_cls_logits, pred_bbox_offsets = self.merge_rpn_score_box( pred_cls_logit_list, pred_bbox_offset_list) fg_mask = rpn_labels > 0 valid_mask = rpn_labels >= 0 num_valid = valid_mask.sum() # rpn classification loss loss_rpn_cls = F.loss.binary_cross_entropy( pred_cls_logits[valid_mask], rpn_labels[valid_mask]) # rpn regression loss loss_rpn_bbox = layers.smooth_l1_loss( pred_bbox_offsets[fg_mask], rpn_offsets[fg_mask], self.cfg.rpn_smooth_l1_beta, ).sum() / F.maximum(num_valid, 1) loss_dict = { "loss_rpn_cls": loss_rpn_cls, "loss_rpn_bbox": loss_rpn_bbox } return rpn_rois, loss_dict else: return rpn_rois