def init_param(self, model_config): self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] # self.rpn_batch_size = model_config['rpn_batch_size'] self.num_cls_samples = model_config['num_cls_samples'] self.num_reg_samples = model_config['num_reg_samples'] self.use_focal_loss = model_config['use_focal_loss'] # sampler # self.sampler = HardNegativeSampler({"fg_fraction": 1.0}) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler({"fg_fraction": 1.0}) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder self.use_iou = model_config.get('use_iou')
def init_param(self, model_config): classes = model_config['classes'] self.classes = classes self.n_classes = len(classes) self.class_agnostic = model_config['class_agnostic'] self.pooling_size = model_config['pooling_size'] self.pooling_mode = model_config['pooling_mode'] self.crop_resize_with_max_pool = model_config[ 'crop_resize_with_max_pool'] self.truncated = model_config['truncated'] self.use_self_attention = model_config.get('use_self_attention') self.use_focal_loss = model_config['use_focal_loss'] self.subsample_twice = model_config['subsample_twice'] self.rcnn_batch_size = model_config['rcnn_batch_size'] # some submodule config self.feature_extractor_config = model_config[ 'feature_extractor_config'] self.rpn_config = model_config['rpn_config'] # assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # sampler # self.bbox_sampler = DetectionSampler({'fg_fraction': 0.5}) # self.bbox_sampler = HardNegativeSampler({'fg_fraction': 1}) # self.bbox_sampler = BalancedSampler({'fg_fraction': 1}) # self.iou_sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler({'fg_fraction': 1})
def init_param(self, model_config): self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] self.rpn_batch_size = model_config['rpn_batch_size'] self.use_focal_loss = model_config['use_focal_loss'] self.iou_criterion = model_config['iou_criterion'] self.use_iox = model_config['use_iox'] self.theta = 1.0 self.alpha = 0.6 # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler(model_config['sampler_config']) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = LEDTargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder
def init_param(self, model_config): self.outside_window_filter = model_config.get('outside_window_filter', False) self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] self.rpn_batch_size = model_config['rpn_batch_size'] self.use_focal_loss = model_config['use_focal_loss'] self.clip_boxes = model_config.get('clip_boxes', True) # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler(model_config['sampler_config']) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder self.use_iou = model_config.get('use_iou')
def init_param(self, model_config): classes = model_config['classes'] self.classes = classes self.n_classes = len(classes) self.class_agnostic = model_config['class_agnostic'] self.pooling_size = model_config['pooling_size'] self.pooling_mode = model_config['pooling_mode'] self.crop_resize_with_max_pool = model_config[ 'crop_resize_with_max_pool'] self.truncated = model_config['truncated'] self.theta = 1.0 self.alpha = 0.6 self.use_focal_loss = model_config['use_focal_loss'] self.subsample_twice = model_config['subsample_twice'] self.rcnn_batch_size = model_config['rcnn_batch_size'] self.iou_criterion = model_config['iou_criterion'] self.use_iox = model_config['use_iox'] # self.use_cls_pred = model_config['use_cls_pred'] # some submodule config self.feature_extractor_config = model_config[ 'feature_extractor_config'] self.rpn_config = model_config['rpn_config'] # assigner self.target_assigner = LEDTargetAssigner( model_config['target_assigner_config']) # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) if self.iou_criterion: self.sampler = DetectionSampler(model_config['sampler_config']) else: self.sampler = BalancedSampler(model_config['sampler_config'])
def init_param(self, model_config): self.feature_extractor_config = model_config['feature_extractor_config'] self.multibox_cfg = [3, 3, 3, 3, 3, 3] self.n_classes = len(model_config['classes']) self.sampler = DetectionSampler(model_config['sampler_config']) self.batch_size = model_config['batch_size'] self.use_focal_loss = model_config['use_focal_loss'] # self.multibox_cfg = model_config['multibox_config'] self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # import ipdb # ipdb.set_trace() self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.bbox_coder = self.target_assigner.bbox_coder
def init_param(self, model_config): self.feat_size = model_config['common_feat_size'] self.batch_size = model_config['batch_size'] self.sample_size = model_config['sample_size'] self.pooling_size = model_config['pooling_size'] self.n_classes = model_config['num_classes'] self.use_focal_loss = model_config['use_focal_loss'] self.feature_extractor_config = model_config[ 'feature_extractor_config'] self.voxel_generator = VoxelGenerator( model_config['voxel_generator_config']) self.voxel_generator.init_voxels() self.integral_map_generator = IntegralMapGenerator() self.oft_target_assigner = OFTargetAssigner( model_config['target_assigner_config']) self.target_assigner = TargetAssigner( model_config['eval_target_assigner_config']) self.target_assigner.analyzer.append_gt = False self.sampler = DetectionSampler(model_config['sampler_config']) self.bbox_coder = self.oft_target_assigner.bbox_coder # find the most expensive operators self.profiler = Profiler() # self.multibin = model_config['multibin'] self.num_bins = model_config['num_bins'] self.reg_channels = 3 + 3 + self.num_bins * 4 # score, pos, dim, ang self.rcnn_output_channels = self.n_classes + self.reg_channels self.rpn_output_channels = 2 + 3 + 3 nms_deltas = model_config.get('nms_deltas') if nms_deltas is None: nms_deltas = 1 self.nms_deltas = nms_deltas
class OrgOHEMThreeIoUFasterRCNN(Model): def forward(self, feed_dict): # import ipdb # ipdb.set_trace() prediction_dict = {} # base model base_feat = self.feature_extractor.first_stage_feature( feed_dict['img']) feed_dict.update({'base_feat': base_feat}) # batch_size = base_feat.shape[0] # rpn model prediction_dict.update(self.rpn_model.forward(feed_dict)) # proposals = prediction_dict['proposals_batch'] # shape(N,num_proposals,5) # pre subsample for reduce consume of memory # if self.training: # self.pre_subsample(prediction_dict, feed_dict) rois_batch = prediction_dict['rois_batch'] # note here base_feat (N,C,H,W),rois_batch (N,num_proposals,5) pooled_feat = self.rcnn_pooling(base_feat, rois_batch.view(-1, 5)) # shape(N,C,1,1) pooled_feat = self.feature_extractor.second_stage_feature(pooled_feat) ######################################## # semantic map ######################################## # no necessary for iou pooled_feat_cls = pooled_feat.mean(3).mean(2) rcnn_cls_scores = self.rcnn_cls_pred(pooled_feat_cls) # rcnn_cls_scores = rcnn_cls_scores_map.mean(3).mean(2) # saliency_map = F.softmax(rcnn_cls_scores_map, dim=1) # rcnn_cls_probs = F.softmax(rcnn_cls_scores, dim=1) # shape(N,C) # rcnn_bbox_feat = pooled_feat * saliency_map[:, 1:, :, :] # rcnn_bbox_feat = rcnn_bbox_feat.mean(3).mean(2) # self attention # import ipdb # ipdb.set_trace() # if self.use_self_attention: # channel_attention = self.generate_channel_attention(pooled_feat) # spatial_attention = self.generate_spatial_attention(pooled_feat) # pooled_feat_reg = pooled_feat * channel_attention # pooled_feat_reg = pooled_feat * spatial_attention # pooled_feat_reg = pooled_feat_reg.mean(3).mean(2) # rcnn_bbox_preds = self.rcnn_bbox_pred(pooled_feat_reg) # else: # rcnn_bbox_preds = self.rcnn_bbox_pred(pooled_feat_cls) # shape(N,C) # pooled_feat = pooled_feat.mean(3).mean(2) # rcnn_bbox_preds = self.rcnn_bbox_pred(pooled_feat) # rcnn_cls_scores = self.rcnn_cls_pred(pooled_feat) # rcnn_cls_probs = F.softmax(rcnn_cls_scores, dim=1) rcnn_bbox_preds = self.rcnn_bbox_pred(pooled_feat) rcnn_bbox_preds = rcnn_bbox_preds.view(rcnn_bbox_preds.shape[0] * 4, -1) rcnn_bbox_preds = rcnn_bbox_preds.mean(-1) # import ipdb # ipdb.set_trace() # select min of abs value # rcnn_bbox_preds = torch.abs(rcnn_bbox_preds) # _, min_idx = torch.min(rcnn_bbox_preds, dim=-1) # row_idx = torch.arange(min_idx.numel()).type_as(min_idx) # rcnn_bbox_preds = rcnn_bbox_preds[row_idx, min_idx] rcnn_bbox_preds = rcnn_bbox_preds.view(-1, 4) # rcnn_bbox_preds = rcnn_bbox_preds.max(3)[0].max(2)[0] # prediction_dict['rcnn_cls_probs'] = rcnn_cls_probs prediction_dict['rcnn_bbox_preds'] = rcnn_bbox_preds prediction_dict['rcnn_cls_scores'] = rcnn_cls_scores # used for track proposals_order = prediction_dict['proposals_order'] prediction_dict['second_rpn_anchors'] = prediction_dict['anchors'][0][ proposals_order] return prediction_dict # def unloaded_parameters(self): # return ['rcnn_cls_pred.bias', 'rcnn_cls_pred.weight'] def generate_channel_attention(self, feat): return feat.mean(3, keepdim=True).mean(2, keepdim=True) def generate_spatial_attention(self, feat): return self.spatial_attention(feat) def init_weights(self): # submodule init weights self.feature_extractor.init_weights() self.rpn_model.init_weights() Filler.normal_init(self.rcnn_cls_pred, 0, 0.01, self.truncated) Filler.normal_init(self.rcnn_bbox_pred, 0, 0.001, self.truncated) def init_modules(self): self.feature_extractor = ResNetFeatureExtractor( self.feature_extractor_config) self.rpn_model = RPNModel(self.rpn_config) self.rcnn_pooling = RoIAlignAvg(self.pooling_size, self.pooling_size, 1.0 / 16.0) # self.rcnn_cls_pred = nn.Conv2d(2048, self.n_classes, 3, 1, 1) self.rcnn_cls_pred = nn.Linear(2048, self.n_classes) self.rcnn_bbox_pred = nn.Conv2d(2048, 4, 3, 1, 1) # if self.class_agnostic: # self.rcnn_bbox_pred = nn.Linear(2048, 4) # else: # self.rcnn_bbox_pred = nn.Linear(2048, 4 * self.n_classes) # loss module # if self.use_focal_loss: # self.rcnn_cls_loss = FocalLoss(2) # else: # self.rcnn_cls_loss = functools.partial( # F.cross_entropy, reduce=False) self.rcnn_cls_loss = nn.MSELoss(reduce=False) self.rcnn_bbox_loss = nn.modules.SmoothL1Loss(reduce=False) # attention self.spatial_attention = nn.Conv2d(2048, 1, 3, 1, 1) def init_param(self, model_config): classes = model_config['classes'] self.classes = classes self.n_classes = len(classes) self.class_agnostic = model_config['class_agnostic'] self.pooling_size = model_config['pooling_size'] self.pooling_mode = model_config['pooling_mode'] self.crop_resize_with_max_pool = model_config[ 'crop_resize_with_max_pool'] self.truncated = model_config['truncated'] self.use_self_attention = model_config.get('use_self_attention') self.use_focal_loss = model_config['use_focal_loss'] self.subsample_twice = model_config['subsample_twice'] self.rcnn_batch_size = model_config['rcnn_batch_size'] # some submodule config self.feature_extractor_config = model_config[ 'feature_extractor_config'] self.rpn_config = model_config['rpn_config'] # assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # sampler # self.bbox_sampler = DetectionSampler({'fg_fraction': 0.5}) # self.bbox_sampler = HardNegativeSampler({'fg_fraction': 1}) # self.bbox_sampler = BalancedSampler({'fg_fraction': 1}) # self.iou_sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler({'fg_fraction': 1}) def loss(self, prediction_dict, feed_dict): """ assign proposals label and subsample from them Then calculate loss """ # import ipdb # ipdb.set_trace() loss_dict = {} # submodule loss loss_dict.update(self.rpn_model.loss(prediction_dict, feed_dict)) rois_batch = prediction_dict['rois_batch'] gt_boxes = feed_dict['gt_boxes'] gt_labels = feed_dict['gt_labels'] ########################## # assigner ########################## rcnn_cls_targets, rcnn_reg_targets, rcnn_cls_weights, rcnn_reg_weights = self.target_assigner.assign( rois_batch[:, :, 1:], gt_boxes, gt_labels) # bounding box regression L1 loss rcnn_bbox_preds = prediction_dict['rcnn_bbox_preds'] rcnn_bbox_loss = self.rcnn_bbox_loss(rcnn_bbox_preds, rcnn_reg_targets[0]).sum(dim=-1) reg_criterion = rcnn_reg_weights * rcnn_bbox_loss # reg_criterion = self.target_assigner.matcher.assigned_overlaps_batch # bbox subsample pos_indicator = rcnn_reg_weights > 0 bbox_batch_sampled_mask = self.sampler.subsample_batch( self.rcnn_batch_size, pos_indicator, criterion=reg_criterion) rcnn_reg_weights *= bbox_batch_sampled_mask.type_as(rcnn_reg_weights) num_reg_coeff = (rcnn_reg_weights > 0).sum(dim=-1) assert num_reg_coeff, 'bug happens' rcnn_bbox_loss *= rcnn_reg_weights[0] rcnn_bbox_loss = rcnn_bbox_loss.sum(dim=-1) / num_reg_coeff.float() # classification loss # rcnn_cls_scores = prediction_dict['rcnn_cls_probs'][:, 1] # rcnn_cls_loss = self.rcnn_cls_loss(rcnn_cls_scores, # rcnn_cls_targets[0]) # cls subsample # cls_criterion = rcnn_cls_loss * rcnn_cls_weights # indicator = rcnn_cls_weights > 0 # pos_indicator = indicator # cls_batch_sampled_mask = self.sampler.subsample_batch( # self.rcnn_batch_size, # pos_indicator, # criterion=cls_criterion, # indicator=indicator) # cls_batch_sampled_mask |= rcnn_reg_weights.type_as( # cls_batch_sampled_mask) # rcnn_cls_weights *= cls_batch_sampled_mask.type_as(rcnn_cls_weights) # num_cls_coeff = (rcnn_cls_weights > 0).sum(dim=-1) # assert num_cls_coeff, 'bug happens' # rcnn_cls_loss *= rcnn_cls_weights[0] # rcnn_cls_loss = rcnn_cls_loss.sum(dim=-1) / num_cls_coeff.float() # loss weights has no gradients # loss_dict['rcnn_cls_loss'] = rcnn_cls_loss loss_dict['rcnn_bbox_loss'] = rcnn_bbox_loss prediction_dict['rcnn_reg_weights'] = rcnn_reg_weights return loss_dict
class RPNModel(Model): def init_param(self, model_config): self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] # self.rpn_batch_size = model_config['rpn_batch_size'] self.num_cls_samples = model_config['num_cls_samples'] self.num_reg_samples = model_config['num_reg_samples'] self.use_focal_loss = model_config['use_focal_loss'] # sampler # self.sampler = HardNegativeSampler({"fg_fraction": 1.0}) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler({"fg_fraction": 1.0}) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder self.use_iou = model_config.get('use_iou') def init_weights(self): self.truncated = False Filler.normal_init(self.rpn_conv, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_cls_score, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_bbox_pred, 0, 0.01, self.truncated) def unfreeze_modules(self): unfreeze_modules = [ self.rpn_coarse_map_conv_iod.bias, self.rpn_fine_map_conv_iod.bias, self.rpn_coarse_map_conv_iog.bias, self.rpn_fine_map_conv_iog.bias, self.rpn_coarse_map_conv_iou.bias, self.rpn_fine_map_conv_iou.bias, self.rpn_coarse_map_conv_iod.weight, self.rpn_fine_map_conv_iod.weight, self.rpn_coarse_map_conv_iog.weight, self.rpn_fine_map_conv_iog.weight, self.rpn_coarse_map_conv_iou.weight, self.rpn_fine_map_conv_iou.weight ] for module in unfreeze_modules: module.requires_grad = True def init_modules(self): # define the convrelu layers processing input feature map self.rpn_conv = nn.Conv2d(self.in_channels, 512, 3, 1, 1, bias=True) # define bg/fg classifcation score layer self.rpn_cls_score = nn.Conv2d(512, self.nc_score_out, 1, 1, 0) # define anchor box offset prediction layer if self.use_score: bbox_feat_channels = 512 + 2 self.nc_bbox_out /= self.num_anchors else: bbox_feat_channels = 512 self.rpn_bbox_pred = nn.Conv2d(bbox_feat_channels, self.nc_bbox_out, 1, 1, 0) # bbox self.rpn_bbox_loss = nn.modules.loss.SmoothL1Loss(reduce=False) # cls if self.use_focal_loss: self.rpn_cls_loss = FocalLoss(2) else: self.rpn_cls_loss = functools.partial(F.cross_entropy, reduce=False) def generate_proposal(self, rpn_cls_probs, anchors, rpn_bbox_preds, im_info): # TODO create a new Function """ Args: rpn_cls_probs: FloatTensor,shape(N,2*num_anchors,H,W) rpn_bbox_preds: FloatTensor,shape(N,num_anchors*4,H,W) anchors: FloatTensor,shape(N,4,H,W) Returns: proposals_batch: FloatTensor, shape(N,post_nms_topN,4) fg_probs_batch: FloatTensor, shape(N,post_nms_topN) """ # assert len( # rpn_bbox_preds) == 1, 'just one feature maps is supported now' # rpn_bbox_preds = rpn_bbox_preds[0] anchors = anchors[0] # do not backward anchors = anchors rpn_cls_probs = rpn_cls_probs.detach() rpn_bbox_preds = rpn_bbox_preds.detach() batch_size = rpn_bbox_preds.shape[0] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(batch_size, -1, 4) # apply deltas to anchors to decode # loop here due to many features maps # proposals = [] # for rpn_bbox_preds_single_map, anchors_single_map in zip( # rpn_bbox_preds, anchors): # proposals.append( # self.bbox_coder.decode(rpn_bbox_preds_single_map, # anchors_single_map)) # proposals = torch.cat(proposals, dim=1) proposals = self.bbox_coder.decode_batch(rpn_bbox_preds, anchors) # filer and clip proposals = box_ops.clip_boxes(proposals, im_info) # fg prob fg_probs = rpn_cls_probs[:, self.num_anchors:, :, :] fg_probs = fg_probs.permute(0, 2, 3, 1).contiguous().view(batch_size, -1) # sort fg _, fg_probs_order = torch.sort(fg_probs, dim=1, descending=True) # fg_probs_batch = torch.zeros(batch_size, # self.post_nms_topN).type_as(rpn_cls_probs) proposals_batch = torch.zeros(batch_size, self.post_nms_topN, 4).type_as(rpn_bbox_preds) proposals_order = torch.zeros( batch_size, self.post_nms_topN).fill_(-1).type_as(fg_probs_order) for i in range(batch_size): proposals_single = proposals[i] fg_probs_single = fg_probs[i] fg_order_single = fg_probs_order[i] # pre nms if self.pre_nms_topN > 0: fg_order_single = fg_order_single[:self.pre_nms_topN] proposals_single = proposals_single[fg_order_single] fg_probs_single = fg_probs_single[fg_order_single] # nms keep_idx_i = nms( torch.cat((proposals_single, fg_probs_single.unsqueeze(1)), 1), self.nms_thresh) keep_idx_i = keep_idx_i.long().view(-1) # post nms if self.post_nms_topN > 0: keep_idx_i = keep_idx_i[:self.post_nms_topN] proposals_single = proposals_single[keep_idx_i, :] fg_probs_single = fg_probs_single[keep_idx_i] fg_order_single = fg_order_single[keep_idx_i] # padding 0 at the end. num_proposal = keep_idx_i.numel() proposals_batch[i, :num_proposal, :] = proposals_single # fg_probs_batch[i, :num_proposal] = fg_probs_single proposals_order[i, :num_proposal] = fg_order_single return proposals_batch, proposals_order def forward(self, bottom_blobs): base_feat = bottom_blobs['base_feat'] batch_size = base_feat.shape[0] gt_boxes = bottom_blobs['gt_boxes'] im_info = bottom_blobs['im_info'] # rpn conv rpn_conv = F.relu(self.rpn_conv(base_feat), inplace=True) # rpn cls score # shape(N,2*num_anchors,H,W) rpn_cls_scores = self.rpn_cls_score(rpn_conv) # rpn cls prob shape(N,2*num_anchors,H,W) rpn_cls_score_reshape = rpn_cls_scores.view(batch_size, 2, -1) rpn_cls_probs = F.softmax(rpn_cls_score_reshape, dim=1) rpn_cls_probs = rpn_cls_probs.view_as(rpn_cls_scores) # import ipdb # ipdb.set_trace() # rpn bbox pred # shape(N,4*num_anchors,H,W) if self.use_score: # shape (N,2,num_anchoros*H*W) rpn_cls_scores = rpn_cls_score_reshape.permute(0, 2, 1) rpn_bbox_preds = [] for i in range(self.num_anchors): rpn_bbox_feat = torch.cat( [rpn_conv, rpn_cls_scores[:, ::self.num_anchors, :, :]], dim=1) rpn_bbox_preds.append(self.rpn_bbox_pred(rpn_bbox_feat)) rpn_bbox_preds = torch.cat(rpn_bbox_preds, dim=1) else: # get rpn offsets to the anchor boxes rpn_bbox_preds = self.rpn_bbox_pred(rpn_conv) # rpn_bbox_preds = [rpn_bbox_preds] # generate anchors feature_map_list = [base_feat.size()[-2:]] anchors = self.anchor_generator.generate(feature_map_list) ############################### # Proposal ############################### # note that proposals_order is used for track transform of propsoals proposals_batch, proposals_order = self.generate_proposal( rpn_cls_probs, anchors, rpn_bbox_preds, im_info) batch_idx = torch.arange(batch_size).view(batch_size, 1).expand( -1, proposals_batch.shape[1]).type_as(proposals_batch) rois_batch = torch.cat((batch_idx.unsqueeze(-1), proposals_batch), dim=2) if self.training: rois_batch = self.append_gt(rois_batch, gt_boxes) rpn_cls_scores = rpn_cls_scores.view(batch_size, 2, -1, rpn_cls_scores.shape[2], rpn_cls_scores.shape[3]) rpn_cls_scores = rpn_cls_scores.permute(0, 3, 4, 2, 1).contiguous().view( batch_size, -1, 2) # postprocess rpn_cls_probs = rpn_cls_probs.view(batch_size, 2, -1, rpn_cls_probs.shape[2], rpn_cls_probs.shape[3]) rpn_cls_probs = rpn_cls_probs.permute(0, 3, 4, 2, 1).contiguous().view( batch_size, -1, 2) predict_dict = { 'proposals_batch': proposals_batch, 'rpn_cls_scores': rpn_cls_scores, 'rois_batch': rois_batch, 'anchors': anchors, # used for loss 'rpn_bbox_preds': rpn_bbox_preds, 'rpn_cls_probs': rpn_cls_probs, 'proposals_order': proposals_order, } return predict_dict def append_gt(self, rois_batch, gt_boxes): ################################ # append gt_boxes to rois_batch for losses ################################ # may be some bugs here gt_boxes_append = torch.zeros(gt_boxes.shape[0], gt_boxes.shape[1], 5).type_as(gt_boxes) gt_boxes_append[:, :, 1:5] = gt_boxes[:, :, :4] # cat gt_boxes to rois_batch rois_batch = torch.cat([rois_batch, gt_boxes_append], dim=1) return rois_batch def loss(self, prediction_dict, feed_dict): # loss for cls loss_dict = {} gt_boxes = feed_dict['gt_boxes'] anchors = prediction_dict['anchors'] assert len(anchors) == 1, 'just one feature maps is supported now' anchors = anchors[0] ################################# # target assigner ################################ # no need gt labels here,it just a binary classifcation problem # import ipdb # ipdb.set_trace() rpn_cls_targets, rpn_reg_targets, \ rpn_cls_weights, rpn_reg_weights = \ self.target_assigner.assign(anchors, gt_boxes, gt_labels=None) ################################ # double subsample ################################ rpn_cls_probs = prediction_dict['rpn_cls_probs'][:, :, 1] cls_criterion = rpn_cls_probs # cls loss rpn_cls_score = prediction_dict['rpn_cls_scores'] # rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score, rpn_cls_targets) rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score.view(-1, 2), rpn_cls_targets.view(-1)) rpn_cls_loss = rpn_cls_loss.view_as(rpn_cls_weights) # cls_criterion = rpn_cls_loss # cls subsample # pos_indicator = rpn_cls_targets > 0 # ignore fg/bg indicator = rpn_cls_weights > 0 pos_indicator = indicator cls_batch_sampled_mask = self.sampler.subsample_batch( self.num_cls_samples, pos_indicator, criterion=cls_criterion, indicator=indicator) cls_batch_sampled_mask = cls_batch_sampled_mask.type_as( rpn_cls_weights) rpn_cls_weights = rpn_cls_weights * cls_batch_sampled_mask num_cls_coeff = (rpn_cls_weights > 0).sum(dim=1) # reg subsample # subsample all from fg pos_indicator = rpn_reg_weights > 0 rpn_bbox_preds = prediction_dict['rpn_bbox_preds'] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(rpn_bbox_preds.shape[0], -1, 4) rpn_reg_loss = self.rpn_bbox_loss(rpn_bbox_preds, rpn_reg_targets) # reg_criterion = rpn_reg_loss.sum(dim=-1) reg_batch_sampled_mask = self.sampler.subsample_batch( self.num_reg_samples, pos_indicator, criterion=cls_criterion, indicator=None) reg_batch_sampled_mask |= cls_batch_sampled_mask.type_as( reg_batch_sampled_mask) reg_batch_sampled_mask = reg_batch_sampled_mask.type_as( rpn_cls_weights) rpn_reg_weights = rpn_reg_weights * reg_batch_sampled_mask num_reg_coeff = (rpn_reg_weights > 0).sum(dim=1) # check # assert num_cls_coeff, 'bug happens' # assert num_reg_coeff, 'bug happens' if num_cls_coeff == 0: num_cls_coeff = torch.ones([]).type_as(num_cls_coeff) if num_reg_coeff == 0: num_reg_coeff = torch.ones([]).type_as(num_reg_coeff) # cls loss # rpn_cls_score = prediction_dict['rpn_cls_scores'] # # rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score, rpn_cls_targets) # rpn_cls_loss = self.rpn_cls_loss( # rpn_cls_score.view(-1, 2), rpn_cls_targets.view(-1)) rpn_cls_loss *= rpn_cls_weights rpn_cls_loss = rpn_cls_loss.sum(dim=1) / num_cls_coeff.float() # bbox loss # shape(N,num,4) # rpn_bbox_preds = prediction_dict['rpn_bbox_preds'] # rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # # shape(N,H*W*num_anchors,4) # rpn_bbox_preds = rpn_bbox_preds.view(rpn_bbox_preds.shape[0], -1, 4) # rpn_reg_loss = self.rpn_bbox_loss(rpn_bbox_preds, rpn_reg_targets) rpn_reg_loss *= rpn_reg_weights.unsqueeze(-1).expand(-1, -1, 4) rpn_reg_loss = rpn_reg_loss.view( rpn_reg_loss.shape[0], -1).sum(dim=1) / num_reg_coeff.float() loss_dict['rpn_cls_loss'] = rpn_cls_loss loss_dict['rpn_bbox_loss'] = rpn_reg_loss return loss_dict
class GateFasterRCNN(Model): def forward(self, feed_dict): # import ipdb # ipdb.set_trace() prediction_dict = {} # base model base_feat = self.feature_extractor.first_stage_feature( feed_dict['img']) feed_dict.update({'base_feat': base_feat}) # batch_size = base_feat.shape[0] # rpn model prediction_dict.update(self.rpn_model.forward(feed_dict)) # proposals = prediction_dict['proposals_batch'] # shape(N,num_proposals,5) # pre subsample for reduce consume of memory if self.training: self.pre_subsample(prediction_dict, feed_dict) rois_batch = prediction_dict['rois_batch'] # note here base_feat (N,C,H,W),rois_batch (N,num_proposals,5) pooled_feat = self.rcnn_pooling(base_feat, rois_batch.view(-1, 5)) # shape(N,C,1,1) pooled_feat = self.feature_extractor.second_stage_feature(pooled_feat) # shape(N,C) pooled_feat = pooled_feat.mean(3).mean(2) rcnn_bbox_preds = self.rcnn_bbox_pred(pooled_feat) rcnn_cls_scores = self.rcnn_cls_pred(pooled_feat) rcnn_cls_probs = F.softmax(rcnn_cls_scores, dim=1) prediction_dict['rcnn_cls_probs'] = rcnn_cls_probs prediction_dict['rcnn_bbox_preds'] = rcnn_bbox_preds prediction_dict['rcnn_cls_scores'] = rcnn_cls_scores # used for track proposals_order = prediction_dict['proposals_order'] prediction_dict['second_rpn_anchors'] = prediction_dict['anchors'][0][ proposals_order] return prediction_dict def init_weights(self): # submodule init weights self.feature_extractor.init_weights() self.rpn_model.init_weights() Filler.normal_init(self.rcnn_cls_pred, 0, 0.01, self.truncated) Filler.normal_init(self.rcnn_bbox_pred, 0, 0.001, self.truncated) def init_modules(self): self.feature_extractor = FeatureExtractor( self.feature_extractor_config) self.rpn_model = GateRPNModel(self.rpn_config) self.rcnn_pooling = RoIAlignAvg(self.pooling_size, self.pooling_size, 1.0 / 16.0) self.rcnn_cls_pred = nn.Linear(2048, self.n_classes) if self.class_agnostic: self.rcnn_bbox_pred = nn.Linear(2048, 4) else: self.rcnn_bbox_pred = nn.Linear(2048, 4 * self.n_classes) # loss module if self.use_focal_loss: self.rcnn_cls_loss = FocalLoss(2) else: self.rcnn_cls_loss = functools.partial(F.cross_entropy, reduce=False) self.rcnn_bbox_loss = nn.modules.SmoothL1Loss(reduce=False) def init_param(self, model_config): classes = model_config['classes'] self.classes = classes self.n_classes = len(classes) self.class_agnostic = model_config['class_agnostic'] self.pooling_size = model_config['pooling_size'] self.pooling_mode = model_config['pooling_mode'] self.crop_resize_with_max_pool = model_config[ 'crop_resize_with_max_pool'] self.truncated = model_config['truncated'] self.use_focal_loss = model_config['use_focal_loss'] self.subsample_twice = model_config['subsample_twice'] self.rcnn_batch_size = model_config['rcnn_batch_size'] # some submodule config self.feature_extractor_config = model_config[ 'feature_extractor_config'] self.rpn_config = model_config['rpn_config'] # assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler(model_config['sampler_config']) def pre_subsample(self, prediction_dict, feed_dict): rois_batch = prediction_dict['rois_batch'] gt_boxes = feed_dict['gt_boxes'] gt_labels = feed_dict['gt_labels'] ########################## # assigner ########################## # import ipdb # ipdb.set_trace() rcnn_cls_targets, rcnn_reg_targets, rcnn_cls_weights, rcnn_reg_weights = self.target_assigner.assign( rois_batch[:, :, 1:], gt_boxes, gt_labels) ########################## # subsampler ########################## pos_indicator = rcnn_cls_targets > 0 indicator = rcnn_cls_weights > 0 # subsample from all # shape (N,M) # use overlaps to subsample use_iou_for_criteron = True if use_iou_for_criteron: cls_criterion = self.target_assigner.matcher.assigned_overlaps_batch else: cls_criterion = None batch_sampled_mask = self.sampler.subsample_batch( self.rcnn_batch_size, pos_indicator, indicator=indicator, criterion=cls_criterion) rcnn_cls_weights = rcnn_cls_weights[batch_sampled_mask] rcnn_reg_weights = rcnn_reg_weights[batch_sampled_mask] num_cls_coeff = rcnn_cls_weights.type( torch.cuda.ByteTensor).sum(dim=-1) num_reg_coeff = rcnn_reg_weights.type( torch.cuda.ByteTensor).sum(dim=-1) # check assert num_cls_coeff, 'bug happens' assert num_reg_coeff, 'bug happens' prediction_dict[ 'rcnn_cls_weights'] = rcnn_cls_weights / num_cls_coeff.float() prediction_dict[ 'rcnn_reg_weights'] = rcnn_reg_weights / num_reg_coeff.float() prediction_dict['rcnn_cls_targets'] = rcnn_cls_targets[ batch_sampled_mask] prediction_dict['rcnn_reg_targets'] = rcnn_reg_targets[ batch_sampled_mask] # update rois_batch prediction_dict['rois_batch'] = rois_batch[batch_sampled_mask].view( rois_batch.shape[0], -1, 5) if not self.training: # used for track proposals_order = prediction_dict['proposals_order'] prediction_dict['proposals_order'] = proposals_order[ batch_sampled_mask] def loss(self, prediction_dict, feed_dict): """ assign proposals label and subsample from them Then calculate loss """ loss_dict = {} # submodule loss loss_dict.update(self.rpn_model.loss(prediction_dict, feed_dict)) # targets and weights rcnn_cls_weights = prediction_dict['rcnn_cls_weights'] rcnn_reg_weights = prediction_dict['rcnn_reg_weights'] rcnn_cls_targets = prediction_dict['rcnn_cls_targets'] rcnn_reg_targets = prediction_dict['rcnn_reg_targets'] # classification loss rcnn_cls_scores = prediction_dict['rcnn_cls_scores'] rcnn_cls_loss = self.rcnn_cls_loss(rcnn_cls_scores, rcnn_cls_targets) rcnn_cls_loss *= rcnn_cls_weights rcnn_cls_loss = rcnn_cls_loss.sum(dim=-1) # bounding box regression L1 loss rcnn_bbox_preds = prediction_dict['rcnn_bbox_preds'] rcnn_bbox_loss = self.rcnn_bbox_loss(rcnn_bbox_preds, rcnn_reg_targets).sum(dim=-1) rcnn_bbox_loss *= rcnn_reg_weights # rcnn_bbox_loss *= rcnn_reg_weights rcnn_bbox_loss = rcnn_bbox_loss.sum(dim=-1) # loss weights has no gradients loss_dict['rcnn_cls_loss'] = rcnn_cls_loss loss_dict['rcnn_bbox_loss'] = rcnn_bbox_loss # add rcnn_cls_targets to get the statics of rpn loss_dict['rcnn_cls_targets'] = rcnn_cls_targets return loss_dict
class RPNModel(Model): def init_param(self, model_config): self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] self.rpn_batch_size = model_config['rpn_batch_size'] self.use_focal_loss = model_config['use_focal_loss'] # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler(model_config['sampler_config']) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder self.use_iou = model_config.get('use_iou') def init_weights(self): self.truncated = False Filler.normal_init(self.rpn_conv, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_cls_score, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_bbox_pred, 0, 0.01, self.truncated) def init_modules(self): # define the convrelu layers processing input feature map self.rpn_conv = nn.Conv2d(self.in_channels, 512, 3, 1, 1, bias=True) # define bg/fg classifcation score layer self.rpn_cls_score = nn.Conv2d(512, self.nc_score_out, 1, 1, 0) # define anchor box offset prediction layer if self.use_score: bbox_feat_channels = 512 + 2 self.nc_bbox_out /= self.num_anchors else: bbox_feat_channels = 512 self.rpn_bbox_pred = nn.Conv2d(bbox_feat_channels, self.nc_bbox_out, 1, 1, 0) # bbox self.rpn_bbox_loss = nn.modules.loss.SmoothL1Loss(reduce=False) # cls if self.use_focal_loss: self.rpn_cls_loss = FocalLoss(2) else: self.rpn_cls_loss = functools.partial(F.cross_entropy, reduce=False) # def generate_proposal(self, rpn_cls_probs, anchors, rpn_bbox_preds, # im_info): # pass def forward(self, bottom_blobs): base_feat = bottom_blobs['base_feat'] batch_size = base_feat.shape[0] gt_boxes = bottom_blobs['gt_boxes'] im_info = bottom_blobs['im_info'] # rpn conv rpn_conv = F.relu(self.rpn_conv(base_feat), inplace=True) # rpn cls score # shape(N,2*num_anchors,H,W) rpn_cls_scores = self.rpn_cls_score(rpn_conv) # rpn cls prob shape(N,2*num_anchors,H,W) rpn_cls_score_reshape = rpn_cls_scores.view(batch_size, 2, -1) rpn_cls_probs = F.softmax(rpn_cls_score_reshape, dim=1) rpn_cls_probs = rpn_cls_probs.view_as(rpn_cls_scores) # import ipdb # ipdb.set_trace() # rpn bbox pred # shape(N,4*num_anchors,H,W) if self.use_score: # shape (N,2,num_anchoros*H*W) rpn_cls_scores = rpn_cls_score_reshape.permute(0, 2, 1) rpn_bbox_preds = [] for i in range(self.num_anchors): rpn_bbox_feat = torch.cat( [rpn_conv, rpn_cls_scores[:, ::self.num_anchors, :, :]], dim=1) rpn_bbox_preds.append(self.rpn_bbox_pred(rpn_bbox_feat)) rpn_bbox_preds = torch.cat(rpn_bbox_preds, dim=1) else: # get rpn offsets to the anchor boxes rpn_bbox_preds = self.rpn_bbox_pred(rpn_conv) # rpn_bbox_preds = [rpn_bbox_preds] # generate anchors feature_map_list = [base_feat.size()[-2:]] anchors = self.anchor_generator.generate(feature_map_list) ############################### # Proposal ############################### # note that proposals_order is used for track transform of propsoals rois_batch, proposals_order = Proposal.apply(rpn_cls_probs, anchors, rpn_bbox_preds, im_info) # batch_idx = torch.arange(batch_size).view(batch_size, 1).expand( # -1, proposals_batch.shape[1]).type_as(proposals_batch) # rois_batch = torch.cat((batch_idx.unsqueeze(-1), proposals_batch), # dim=2) if self.training: rois_batch = self.append_gt(rois_batch, gt_boxes) rpn_cls_scores = rpn_cls_scores.view(batch_size, 2, -1, rpn_cls_scores.shape[2], rpn_cls_scores.shape[3]) rpn_cls_scores = rpn_cls_scores.permute(0, 3, 4, 2, 1).contiguous().view( batch_size, -1, 2) # postprocess rpn_cls_probs = rpn_cls_probs.view(batch_size, 2, -1, rpn_cls_probs.shape[2], rpn_cls_probs.shape[3]) rpn_cls_probs = rpn_cls_probs.permute(0, 3, 4, 2, 1).contiguous().view( batch_size, -1, 2) predict_dict = { 'rpn_cls_scores': rpn_cls_scores, 'rois_batch': rois_batch, 'anchors': anchors, # used for loss 'rpn_bbox_preds': rpn_bbox_preds, 'rpn_cls_probs': rpn_cls_probs, 'proposals_order': proposals_order, } return predict_dict def append_gt(self, rois_batch, gt_boxes): ################################ # append gt_boxes to rois_batch for losses ################################ # may be some bugs here gt_boxes_append = torch.zeros(gt_boxes.shape[0], gt_boxes.shape[1], 5).type_as(gt_boxes) gt_boxes_append[:, :, 1:5] = gt_boxes[:, :, :4] # cat gt_boxes to rois_batch rois_batch = torch.cat([rois_batch, gt_boxes_append], dim=1) return rois_batch def loss(self, prediction_dict, feed_dict): # loss for cls loss_dict = {} gt_boxes = feed_dict['gt_boxes'] anchors = prediction_dict['anchors'] assert len(anchors) == 1, 'just one feature maps is supported now' anchors = anchors[0] ################################# # target assigner ################################ # no need gt labels here,it just a binary classifcation problem # import ipdb # ipdb.set_trace() rpn_cls_targets, rpn_reg_targets, \ rpn_cls_weights, rpn_reg_weights = \ self.target_assigner.assign(anchors, gt_boxes, gt_labels=None) ################################ # subsample ################################ pos_indicator = rpn_reg_weights > 0 indicator = rpn_cls_weights > 0 if self.use_iou: cls_criterion = self.target_assigner.matcher.assigned_overlaps_batch else: rpn_cls_probs = prediction_dict['rpn_cls_probs'][:, :, 1] cls_criterion = rpn_cls_probs batch_sampled_mask = self.sampler.subsample_batch( self.rpn_batch_size, pos_indicator, criterion=cls_criterion, indicator=indicator) batch_sampled_mask = batch_sampled_mask.type_as(rpn_cls_weights) rpn_cls_weights = rpn_cls_weights * batch_sampled_mask rpn_reg_weights = rpn_reg_weights * batch_sampled_mask num_cls_coeff = (rpn_cls_weights > 0).sum(dim=1) num_reg_coeff = (rpn_reg_weights > 0).sum(dim=1) # check # assert num_cls_coeff, 'bug happens' # assert num_reg_coeff, 'bug happens' if num_cls_coeff == 0: num_cls_coeff = torch.ones([]).type_as(num_cls_coeff) if num_reg_coeff == 0: num_reg_coeff = torch.ones([]).type_as(num_reg_coeff) # cls loss rpn_cls_score = prediction_dict['rpn_cls_scores'] # rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score, rpn_cls_targets) rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score.view(-1, 2), rpn_cls_targets.view(-1)) rpn_cls_loss = rpn_cls_loss.view_as(rpn_cls_weights) rpn_cls_loss *= rpn_cls_weights rpn_cls_loss = rpn_cls_loss.sum(dim=1) / num_cls_coeff.float() # bbox loss # shape(N,num,4) rpn_bbox_preds = prediction_dict['rpn_bbox_preds'] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(rpn_bbox_preds.shape[0], -1, 4) rpn_reg_loss = self.rpn_bbox_loss(rpn_bbox_preds, rpn_reg_targets) rpn_reg_loss *= rpn_reg_weights.unsqueeze(-1).expand(-1, -1, 4) rpn_reg_loss = rpn_reg_loss.view( rpn_reg_loss.shape[0], -1).sum(dim=1) / num_reg_coeff.float() loss_dict['rpn_cls_loss'] = rpn_cls_loss loss_dict['rpn_bbox_loss'] = rpn_reg_loss return loss_dict
class OrgOHEMThreeIoUSecondStageFasterRCNN(Model): def forward(self, feed_dict): # import ipdb # ipdb.set_trace() prediction_dict = {} # base model base_feat = self.feature_extractor.first_stage_feature( feed_dict['img']) feed_dict.update({'base_feat': base_feat}) # batch_size = base_feat.shape[0] # rpn model prediction_dict.update(self.rpn_model.forward(feed_dict)) rois_batch = prediction_dict['rois_batch'] # note here base_feat (N,C,H,W),rois_batch (N,num_proposals,5) pooled_feat = self.rcnn_pooling(base_feat, rois_batch.view(-1, 5)) # shape(N,C,1,1) pooled_feat_reg = self.feature_extractor.second_stage_feature( pooled_feat) ######################################## # semantic map ######################################## # no necessary for iou # pooled_feat_reg = pooled_feat_reg.mean(3).mean(2) pooled_feat_cls = self.feature_extractor.third_stage_feature( pooled_feat) pooled_feat_cls = pooled_feat_cls.mean(3).mean(2) rcnn_cls_scores = self.rcnn_cls_pred(pooled_feat_cls) # rcnn_cls_scores = rcnn_cls_scores_map.mean(3).mean(2) # saliency_map = F.softmax(rcnn_cls_scores_map, dim=1) rcnn_cls_probs = F.softmax(rcnn_cls_scores, dim=1) pooled_feat_reg = pooled_feat_reg.mean(3).mean(2) rcnn_bbox_preds = self.rcnn_bbox_pred(pooled_feat_reg) prediction_dict['rcnn_cls_probs'] = rcnn_cls_probs prediction_dict['rcnn_bbox_preds'] = rcnn_bbox_preds prediction_dict['rcnn_cls_scores'] = rcnn_cls_scores # used for track proposals_order = prediction_dict['proposals_order'] prediction_dict['second_rpn_anchors'] = prediction_dict['anchors'][0][ proposals_order] return prediction_dict def unfreeze_part_modules(self): model = self.feature_extractor.third_stage_feature for param in model.parameters(): param.requires_grad = True def init_weights(self): # submodule init weights self.feature_extractor.init_weights() self.rpn_model.init_weights() Filler.normal_init(self.rcnn_cls_pred, 0, 0.01, self.truncated) Filler.normal_init(self.rcnn_bbox_pred, 0, 0.001, self.truncated) self.freeze_modules() self.unfreeze_part_modules() def init_modules(self): self.feature_extractor = ResNetFeatureExtractor( self.feature_extractor_config) self.rpn_model = RPNModel(self.rpn_config) self.rcnn_pooling = RoIAlignAvg(self.pooling_size, self.pooling_size, 1.0 / 16.0) # self.rcnn_cls_pred = nn.Conv2d(2048, self.n_classes, 3, 1, 1) # import ipdb # ipdb.set_trace() self.rcnn_cls_pred = nn.Linear(2048, self.n_classes) # self.rcnn_bbox_pred = nn.Conv2d(2048, 4, 3, 1, 1) # if self.class_agnostic: self.rcnn_bbox_pred = nn.Linear(2048, 4) # else: # self.rcnn_bbox_pred = nn.Linear(2048, 4 * self.n_classes) # loss module # if self.use_focal_loss: # self.rcnn_cls_loss = FocalLoss(2) # else: # self.rcnn_cls_loss = functools.partial( # F.cross_entropy, reduce=False) self.rcnn_cls_loss = nn.MSELoss(reduce=False) self.rcnn_bbox_loss = nn.modules.SmoothL1Loss(reduce=False) # attention self.spatial_attention = nn.Conv2d(2048, 1, 3, 1, 1) def init_param(self, model_config): classes = model_config['classes'] self.classes = classes self.n_classes = len(classes) self.class_agnostic = model_config['class_agnostic'] self.pooling_size = model_config['pooling_size'] self.pooling_mode = model_config['pooling_mode'] self.crop_resize_with_max_pool = model_config[ 'crop_resize_with_max_pool'] self.truncated = model_config['truncated'] self.use_self_attention = model_config.get('use_self_attention') self.use_focal_loss = model_config['use_focal_loss'] self.subsample_twice = model_config['subsample_twice'] self.rcnn_batch_size = model_config['rcnn_batch_size'] # some submodule config self.feature_extractor_config = model_config[ 'feature_extractor_config'] self.rpn_config = model_config['rpn_config'] # assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # sampler # self.bbox_sampler = DetectionSampler({'fg_fraction': 0.5}) # self.bbox_sampler = HardNegativeSampler({'fg_fraction': 1}) # self.bbox_sampler = BalancedSampler({'fg_fraction': 1}) # self.iou_sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler({'fg_fraction': 1}) def loss(self, prediction_dict, feed_dict): """ assign proposals label and subsample from them Then calculate loss """ # import ipdb # ipdb.set_trace() loss_dict = {} # submodule loss loss_dict.update(self.rpn_model.loss(prediction_dict, feed_dict)) rois_batch = prediction_dict['rois_batch'] gt_boxes = feed_dict['gt_boxes'] gt_labels = feed_dict['gt_labels'] ########################## # assigner ########################## rcnn_cls_targets, rcnn_reg_targets, rcnn_cls_weights, rcnn_reg_weights = self.target_assigner.assign( rois_batch[:, :, 1:], gt_boxes, gt_labels) # bounding box regression L1 loss # rcnn_bbox_preds = prediction_dict['rcnn_bbox_preds'] # rcnn_bbox_loss = self.rcnn_bbox_loss(rcnn_bbox_preds, # rcnn_reg_targets[0]).sum(dim=-1) # reg_criterion = rcnn_reg_weights * rcnn_bbox_loss # reg_criterion = self.target_assigner.matcher.assigned_overlaps_batch # bbox subsample # pos_indicator = rcnn_reg_weights > 0 # bbox_batch_sampled_mask = self.sampler.subsample_batch( # self.rcnn_batch_size, pos_indicator, criterion=reg_criterion) # rcnn_reg_weights *= bbox_batch_sampled_mask.type_as(rcnn_reg_weights) # num_reg_coeff = (rcnn_reg_weights > 0).sum(dim=-1) # assert num_reg_coeff, 'bug happens' # rcnn_bbox_loss *= rcnn_reg_weights[0] # rcnn_bbox_loss = rcnn_bbox_loss.sum(dim=-1) / num_reg_coeff.float() # classification loss rcnn_cls_scores = prediction_dict['rcnn_cls_probs'][:, 1] rcnn_cls_loss = self.rcnn_cls_loss(rcnn_cls_scores, rcnn_cls_targets[0]) # cls subsample cls_criterion = rcnn_cls_loss * rcnn_cls_weights indicator = rcnn_cls_weights > 0 pos_indicator = indicator cls_batch_sampled_mask = self.sampler.subsample_batch( self.rcnn_batch_size, pos_indicator, criterion=cls_criterion, indicator=indicator) # cls_batch_sampled_mask |= rcnn_reg_weights.type_as( # cls_batch_sampled_mask) rcnn_cls_weights *= cls_batch_sampled_mask.type_as(rcnn_cls_weights) num_cls_coeff = (rcnn_cls_weights > 0).sum(dim=-1) assert num_cls_coeff, 'bug happens' rcnn_cls_loss *= rcnn_cls_weights[0] rcnn_cls_loss = rcnn_cls_loss.sum(dim=-1) / num_cls_coeff.float() # loss weights has no gradients loss_dict['rcnn_cls_loss'] = rcnn_cls_loss # loss_dict['rcnn_bbox_loss'] = rcnn_bbox_loss # analysis precision # import ipdb # ipdb.set_trace() rcnn_cls_probs = prediction_dict['rcnn_cls_probs'] fake_match = self.target_assigner.analyzer.match num_gt = feed_dict['gt_labels'].numel() self.target_assigner.analyzer.analyze_ap(fake_match, rcnn_cls_probs[:, 1], num_gt, thresh=0.5) prediction_dict['rcnn_reg_weights'] = rcnn_reg_weights return loss_dict
class LEDRPNModel(Model): def init_param(self, model_config): self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] self.rpn_batch_size = model_config['rpn_batch_size'] self.use_focal_loss = model_config['use_focal_loss'] self.alpha = 0.6 self.theta = 1.0 self.iox_bbox_coder = DiscreteBBoxCoder( model_config['iox_coder_config']) self.use_sharpL2 = model_config.get('use_sharpL2') self.use_sigmoid = model_config['use_sigmoid'] self.use_cls_pred = model_config['use_cls_pred'] # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler(model_config['sampler_config']) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = LEDTargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder def iox_clip(self, iox): iox = iox.clone() iox[iox < 0] = 0 iox[iox > 1] = 1 return iox def init_weights(self): self.truncated = False Filler.normal_init(self.rpn_conv, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_cls_score, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_bbox_pred, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_coarse_map_conv_iod, 0, 0.001, self.truncated) Filler.normal_init(self.rpn_fine_map_conv_iod, 0, 0.001, self.truncated) Filler.normal_init(self.rpn_coarse_map_conv_iou, 0, 0.001, self.truncated) Filler.normal_init(self.rpn_fine_map_conv_iou, 0, 0.001, self.truncated) Filler.normal_init(self.rpn_fine_map_conv_iog, 0, 0.001, self.truncated) Filler.normal_init(self.rpn_coarse_map_conv_iog, 0, 0.001, self.truncated) def unfreeze_modules(self): unfreeze_modules = [ self.rpn_coarse_map_conv_iod.bias, self.rpn_fine_map_conv_iod.bias, self.rpn_coarse_map_conv_iog.bias, self.rpn_fine_map_conv_iog.bias, self.rpn_coarse_map_conv_iou.bias, self.rpn_fine_map_conv_iou.bias, self.rpn_coarse_map_conv_iod.weight, self.rpn_fine_map_conv_iod.weight, self.rpn_coarse_map_conv_iog.weight, self.rpn_fine_map_conv_iog.weight, self.rpn_coarse_map_conv_iou.weight, self.rpn_fine_map_conv_iou.weight ] for module in unfreeze_modules: module.requires_grad = True def init_modules(self): # define the convrelu layers processing input feature map self.rpn_conv = nn.Conv2d(self.in_channels, 512, 3, 1, 1, bias=True) # define bg/fg classifcation score layer self.rpn_cls_score = nn.Conv2d(512, self.nc_score_out, 1, 1, 0) self.rpn_coarse_map_conv_iou = nn.Conv2d(512, 4 * self.num_anchors, 1, 1, 0) self.rpn_fine_map_conv_iou = nn.Conv2d(512, 4 * self.num_anchors, 1, 1, 0) self.rpn_coarse_map_conv_iog = nn.Conv2d(512, 4 * self.num_anchors, 1, 1, 0) self.rpn_fine_map_conv_iog = nn.Conv2d(512, 4 * self.num_anchors, 1, 1, 0) self.rpn_coarse_map_conv_iod = nn.Conv2d(512, 4 * self.num_anchors, 1, 1, 0) self.rpn_fine_map_conv_iod = nn.Conv2d(512, 4 * self.num_anchors, 1, 1, 0) # define anchor box offset prediction layer if self.use_score: bbox_feat_channels = 512 + 2 self.nc_bbox_out /= self.num_anchors else: bbox_feat_channels = 512 self.rpn_bbox_pred = nn.Conv2d(bbox_feat_channels, self.nc_bbox_out, 1, 1, 0) # rpn bbox self.rpn_bbox_loss = nn.modules.loss.SmoothL1Loss(reduce=False) if self.use_sharpL2: self.reg_loss = SharpL2Loss() else: self.reg_loss = nn.MSELoss(reduce=False) self.cls_loss = nn.CrossEntropyLoss(reduce=False) # rpn cls if self.use_focal_loss: self.rpn_cls_loss = FocalLoss(2) else: self.rpn_cls_loss = functools.partial(F.cross_entropy, reduce=False) def generate_proposal(self, rpn_cls_probs, anchors, rpn_bbox_preds, im_info): # TODO create a new Function """ Args: rpn_cls_probs: FloatTensor,shape(N,2*num_anchors,H,W) rpn_bbox_preds: FloatTensor,shape(N,num_anchors*4,H,W) anchors: FloatTensor,shape(N,4,H,W) Returns: proposals_batch: FloatTensor, shape(N,post_nms_topN,4) fg_probs_batch: FloatTensor, shape(N,post_nms_topN) """ # assert len( # rpn_bbox_preds) == 1, 'just one feature maps is supported now' # rpn_bbox_preds = rpn_bbox_preds[0] anchors = anchors[0] # do not backward anchors = anchors rpn_fg_cls_probs = rpn_cls_probs.detach() rpn_bbox_preds = rpn_bbox_preds.detach() batch_size = rpn_bbox_preds.shape[0] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(batch_size, -1, 4) # apply deltas to anchors to decode # loop here due to many features maps # proposals = [] # for rpn_bbox_preds_single_map, anchors_single_map in zip( # rpn_bbox_preds, anchors): # proposals.append( # self.bbox_coder.decode(rpn_bbox_preds_single_map, # anchors_single_map)) # proposals = torch.cat(proposals, dim=1) proposals = self.bbox_coder.decode_batch(rpn_bbox_preds, anchors) # filer and clip proposals = box_ops.clip_boxes(proposals, im_info) # fg prob # fg_probs = rpn_cls_probs[:, self.num_anchors:, :, :] # fg_probs = fg_probs.permute(0, 2, 3, 1).contiguous().view(batch_size, # -1) fg_probs = rpn_fg_cls_probs # sort fg _, fg_probs_order = torch.sort(fg_probs, dim=1, descending=True) # fg_probs_batch = torch.zeros(batch_size, # self.post_nms_topN).type_as(rpn_cls_probs) proposals_batch = torch.zeros(batch_size, self.post_nms_topN, 4).type_as(rpn_bbox_preds) proposals_order = torch.zeros( batch_size, self.post_nms_topN).fill_(-1).type_as(fg_probs_order) for i in range(batch_size): proposals_single = proposals[i] fg_probs_single = fg_probs[i] fg_order_single = fg_probs_order[i] # pre nms if self.pre_nms_topN > 0: fg_order_single = fg_order_single[:self.pre_nms_topN] proposals_single = proposals_single[fg_order_single] fg_probs_single = fg_probs_single[fg_order_single] # nms keep_idx_i = nms( torch.cat((proposals_single, fg_probs_single.unsqueeze(1)), 1), self.nms_thresh) keep_idx_i = keep_idx_i.long().view(-1) # post nms if self.post_nms_topN > 0: keep_idx_i = keep_idx_i[:self.post_nms_topN] proposals_single = proposals_single[keep_idx_i, :] fg_probs_single = fg_probs_single[keep_idx_i] fg_order_single = fg_order_single[keep_idx_i] # padding 0 at the end. num_proposal = keep_idx_i.numel() proposals_batch[i, :num_proposal, :] = proposals_single # fg_probs_batch[i, :num_proposal] = fg_probs_single proposals_order[i, :num_proposal] = fg_order_single return proposals_batch, proposals_order def iou_pred(self, rpn_conv): return self.iox_pred(rpn_conv, self.rpn_coarse_map_conv_iou, self.rpn_fine_map_conv_iou) def iog_pred(self, rpn_conv): return self.iox_pred(rpn_conv, self.rpn_coarse_map_conv_iog, self.rpn_fine_map_conv_iog) def iod_pred(self, rpn_conv): return self.iox_pred(rpn_conv, self.rpn_coarse_map_conv_iod, self.rpn_fine_map_conv_iod) def iox_pred(self, rpn_conv, rpn_coarse_map_conv, rpn_fine_map_conv): batch_size = rpn_conv.shape[0] coarse_map = rpn_coarse_map_conv(rpn_conv) fine_map = rpn_fine_map_conv(rpn_conv) coarse_map_reshape = coarse_map.view(batch_size, 4, -1) iou_level_probs = F.softmax(coarse_map_reshape, dim=1) iou_level_probs = iou_level_probs.view_as(coarse_map) if self.use_sigmoid: # normalize it iou_reg = 2 * F.sigmoid(fine_map) - 1 else: iou_reg = fine_map # reshape preprocess iou_reg = iou_reg.view(batch_size, 4, self.num_anchors, -1).permute( 0, 3, 2, 1).contiguous().view(batch_size, -1, 4) iou_cls = iou_level_probs.view(batch_size, 4, self.num_anchors, -1).permute(0, 3, 2, 1).contiguous().view( batch_size, -1, 4) decoded_iou = self.iox_bbox_coder.decode_batch(iou_cls, iou_reg) # used for cls and reg loss iou_cls_scores = coarse_map.view(batch_size, 4, self.num_anchors, -1).permute(0, 3, 2, 1).contiguous().view( batch_size, -1, 4) return decoded_iou, iou_cls_scores, iou_reg def calculate_iou(self, iog, iod): mask = ~(iod == 0) iou_indirect = torch.zeros_like(iog) iod = iod[mask] iog = iog[mask] iou_indirect[mask] = (iod * iog) / (iod + iog - iod * iog) return iou_indirect def forward(self, bottom_blobs): # import ipdb # ipdb.set_trace() base_feat = bottom_blobs['base_feat'] batch_size = base_feat.shape[0] gt_boxes = bottom_blobs['gt_boxes'] # im_info = bottom_blobs['im_info'] im_info = bottom_blobs['input_size'] # rpn conv rpn_conv = F.relu(self.rpn_conv(base_feat), inplace=True) # rpn cls rpn_cls_scores = self.rpn_cls_score(rpn_conv) rpn_cls_score_reshape = rpn_cls_scores.view(batch_size, 2, -1) rpn_cls_probs = F.softmax(rpn_cls_score_reshape, dim=1).view_as(rpn_cls_scores) rpn_cls_probs = rpn_cls_probs.view(batch_size, 2, self.num_anchors, -1).permute(0, 3, 2, 1).contiguous().view( batch_size, -1, 2) rpn_cls_scores = rpn_cls_scores.view(batch_size, 2, self.num_anchors, -1).permute(0, 3, 2, 1).contiguous().view( batch_size, -1, 2) iou, iou_scores, iou_reg = self.iou_pred(rpn_conv) iog, iog_scores, iog_reg = self.iog_pred(rpn_conv) iod, iod_scores, iod_reg = self.iod_pred(rpn_conv) # bugs here iou = self.iox_clip(iou) iog = self.iox_clip(iog) iod = self.iox_clip(iod) iou_indirect = self.calculate_iou(iog, iod) iou_final = (1 - self.alpha) * iou_indirect + self.alpha * iou # import ipdb # ipdb.set_trace() rpn_fg_probs_final = rpn_cls_probs[:, :, 1] * torch.exp(-torch.pow( (1 - iou_final), 2) / self.theta) # rpn_cls_score_reshape = rpn_cls_scores.view(batch_size, 2, -1) # rpn_cls_probs = F.softmax(rpn_cls_score_reshape, dim=1) # rpn_cls_probs = rpn_cls_probs.view_as(rpn_cls_scores) # import ipdb # ipdb.set_trace() # rpn bbox pred # shape(N,4*num_anchors,H,W) # if self.use_score: # # shape (N,2,num_anchoros*H*W) # rpn_cls_scores = rpn_cls_score_reshape.permute(0, 2, 1) # rpn_bbox_preds = [] # for i in range(self.num_anchors): # rpn_bbox_feat = torch.cat( # [rpn_conv, rpn_cls_scores[:, ::self.num_anchors, :, :]], # dim=1) # rpn_bbox_preds.append(self.rpn_bbox_pred(rpn_bbox_feat)) # rpn_bbox_preds = torch.cat(rpn_bbox_preds, dim=1) # else: # get rpn offsets to the anchor boxes rpn_bbox_preds = self.rpn_bbox_pred(rpn_conv) # rpn_bbox_preds = [rpn_bbox_preds] # generate anchors feature_map_list = [base_feat.size()[-2:]] anchors = self.anchor_generator.generate(feature_map_list) ############################### # Proposal ############################### # note that proposals_order is used for tracking transform of propsoals proposals_batch, proposals_order = self.generate_proposal( rpn_fg_probs_final, anchors, rpn_bbox_preds, im_info) batch_idx = torch.arange(batch_size).view(batch_size, 1).expand( -1, proposals_batch.shape[1]).type_as(proposals_batch) rois_batch = torch.cat((batch_idx.unsqueeze(-1), proposals_batch), dim=2) if self.training: rois_batch = self.append_gt(rois_batch, gt_boxes) predict_dict = { 'proposals_batch': proposals_batch, # used for sorting 'rpn_iou_final': rpn_fg_probs_final, 'rois_batch': rois_batch, 'anchors': anchors, # used for loss 'rpn_bbox_preds': rpn_bbox_preds, 'proposals_order': proposals_order, # reg 'rpn_iou_reg': iou_reg, 'rpn_iog_reg': iog_reg, 'rpn_iod_reg': iod_reg, # cls 'rpn_iou_scores': iou_scores, 'rpn_iog_scores': iog_scores, 'rpn_iod_scores': iod_scores, 'rpn_cls_scores': rpn_cls_scores } return predict_dict def append_gt(self, rois_batch, gt_boxes): ################################ # append gt_boxes to rois_batch for losses ################################ # may be some bugs here gt_boxes_append = torch.zeros(gt_boxes.shape[0], gt_boxes.shape[1], 5).type_as(gt_boxes) gt_boxes_append[:, :, 1:5] = gt_boxes[:, :, :4] # cat gt_boxes to rois_batch rois_batch = torch.cat([rois_batch, gt_boxes_append], dim=1) return rois_batch def loss(self, prediction_dict, feed_dict): # loss for cls loss_dict = {} gt_boxes = feed_dict['gt_boxes'] anchors = prediction_dict['anchors'] assert len(anchors) == 1, 'just one feature maps is supported now' anchors = anchors[0] ################################# # target assigner ################################ # no need gt labels here,it just a binary classifcation problem rpn_cls_targets, rpn_reg_targets, \ rpn_cls_weights, rpn_reg_weights = \ self.target_assigner.assign(anchors, gt_boxes, gt_labels=None, input_size=feed_dict['input_size']) ################################ # subsample ################################ rpn_cls_probs = prediction_dict['rpn_iou_final'] cls_criterion = rpn_cls_probs pos_indicator = rpn_reg_weights > 0 indicator = rpn_cls_weights > 0 batch_sampled_mask = self.sampler.subsample_batch( self.rpn_batch_size, pos_indicator, criterion=cls_criterion, indicator=indicator) batch_sampled_mask = batch_sampled_mask.type_as(rpn_cls_weights) rpn_cls_weights = rpn_cls_weights * batch_sampled_mask rpn_reg_weights = rpn_reg_weights * batch_sampled_mask num_cls_coeff = (rpn_cls_weights > 0).sum(dim=1) num_reg_coeff = (rpn_reg_weights > 0).sum(dim=1) # check # assert num_cls_coeff, 'bug happens' # assert num_reg_coeff, 'bug happens' if num_cls_coeff == 0: num_cls_coeff = torch.ones([]).type_as(num_cls_coeff) if num_reg_coeff == 0: num_reg_coeff = torch.ones([]).type_as(num_reg_coeff) # iou loss iou_scores = prediction_dict['rpn_iou_scores'] iou = prediction_dict['rpn_iou_reg'] iou_reg_targets = self.target_assigner.matcher.assigned_overlaps_batch iou_reg_targets_encode = self.iox_bbox_coder.encode_reg( iou_reg_targets) iou_reg_loss = self.reg_loss(iou, iou_reg_targets_encode).sum(dim=-1) iou_scores_targets = self.iox_bbox_coder.encode_cls(iou_reg_targets) iou_cls_loss = self.cls_loss(iou_scores.view(-1, 4), iou_scores_targets.view(-1)) iou_cls_loss = iou_cls_loss.view_as(rpn_cls_weights) iou_cls_loss = iou_cls_loss.mean(dim=1) iou_reg_loss = iou_reg_loss.mean(dim=1) # iog loss iog_scores = prediction_dict['rpn_iog_scores'] iog = prediction_dict['rpn_iog_reg'] iog_reg_targets = self.target_assigner.matcher.assigned_iog_batch iog_reg_targets_encode = self.iox_bbox_coder.encode_reg( iog_reg_targets) iog_reg_loss = self.reg_loss(iog, iog_reg_targets_encode).sum(dim=-1) iog_scores_targets = self.iox_bbox_coder.encode_cls(iog_reg_targets) iog_cls_loss = self.cls_loss(iog_scores.view(-1, 4), iog_scores_targets.view(-1)) iog_cls_loss = iog_cls_loss.view_as(rpn_cls_weights) iog_cls_loss = iog_cls_loss.mean(dim=1) iog_reg_loss = iog_reg_loss.mean(dim=1) # iod loss iod_scores = prediction_dict['rpn_iod_scores'] iod = prediction_dict['rpn_iod_reg'] iod_reg_targets = self.target_assigner.matcher.assigned_iod_batch iod_reg_targets_encode = self.iox_bbox_coder.encode_reg( iod_reg_targets) iod_reg_loss = self.reg_loss(iod, iod_reg_targets_encode).sum(dim=-1) iod_scores_targets = self.iox_bbox_coder.encode_cls(iod_reg_targets) iod_cls_loss = self.cls_loss(iod_scores.view(-1, 4), iod_scores_targets.view(-1)) iod_cls_loss = iod_cls_loss.view_as(rpn_cls_weights) iod_cls_loss = iod_cls_loss.mean(dim=1) iod_reg_loss = iod_reg_loss.mean(dim=1) # cls loss if self.use_cls_pred: rpn_cls_score = prediction_dict['rpn_cls_scores'] rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score.view(-1, 2), rpn_cls_targets.view(-1)) rpn_cls_loss = rpn_cls_loss.view_as(rpn_cls_weights) rpn_cls_loss *= rpn_cls_weights rpn_cls_loss = rpn_cls_loss.sum(dim=1) / num_cls_coeff.float() loss_dict['rpn/cls_loss'] = rpn_cls_loss loss_dict['rpn/iou_cls_loss'] = iou_cls_loss loss_dict['rpn/iou_reg_loss'] = iou_reg_loss loss_dict['rpn/iog_cls_loss'] = iog_cls_loss loss_dict['rpn/iog_reg_loss'] = iog_reg_loss loss_dict['rpn/iod_reg_loss'] = iod_reg_loss loss_dict['rpn/iod_cls_loss'] = iod_cls_loss # bbox loss # shape(N,num,4) rpn_bbox_preds = prediction_dict['rpn_bbox_preds'] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(rpn_bbox_preds.shape[0], -1, 4) rpn_reg_loss = self.rpn_bbox_loss(rpn_bbox_preds, rpn_reg_targets) rpn_reg_loss *= rpn_reg_weights.unsqueeze(-1).expand(-1, -1, 4) rpn_reg_loss = rpn_reg_loss.view( rpn_reg_loss.shape[0], -1).sum(dim=1) / num_reg_coeff.float() # loss_dict['rpn_cls_loss'] = iox_loss loss_dict['rpn/bbox_loss'] = rpn_reg_loss # loss_dict['iox_loss'] = iox_loss return loss_dict
class SSDModel(Model): def init_param(self, model_config): self.feature_extractor_config = model_config['feature_extractor_config'] self.multibox_cfg = [3, 3, 3, 3, 3, 3] self.n_classes = len(model_config['classes']) self.sampler = DetectionSampler(model_config['sampler_config']) self.batch_size = model_config['batch_size'] self.use_focal_loss = model_config['use_focal_loss'] # self.multibox_cfg = model_config['multibox_config'] self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # import ipdb # ipdb.set_trace() self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.bbox_coder = self.target_assigner.bbox_coder # self.priorsbox = PriorBox(model_config['anchor_generator_config']) def init_modules(self): self.feature_extractor = PyramidVggnetExtractor( self.feature_extractor_config) # loc layers and conf layers base_feat = self.feature_extractor.base_feat extra_layers = self.feature_extractor.extras_layers loc_layers, conf_layers = self.make_multibox(base_feat, extra_layers) self.loc_layers = loc_layers self.conf_layers = conf_layers # self.rcnn_3d_preds = nn.Linear() # loss layers self.loc_loss = nn.SmoothL1Loss(reduce=False) if self.use_focal_loss: self.conf_loss = FocalLoss( self.n_classes, alpha=0.2, gamma=2, auto_alpha=False) else: self.conf_loss = nn.CrossEntropyLoss(reduce=False) def make_multibox(self, vgg, extra_layers): cfg = self.multibox_cfg num_classes = self.n_classes loc_layers = [] conf_layers = [] vgg_source = [21, -2] for k, v in enumerate(vgg_source): loc_layers += [ nn.Conv2d( vgg[v].out_channels, cfg[k] * 4, kernel_size=3, padding=1) ] conf_layers += [ nn.Conv2d( vgg[v].out_channels, cfg[k] * num_classes, kernel_size=3, padding=1) ] for k, v in enumerate(extra_layers[1::2], 2): loc_layers += [ nn.Conv2d( v.out_channels, cfg[k] * 4, kernel_size=3, padding=1) ] conf_layers += [ nn.Conv2d( v.out_channels, cfg[k] * num_classes, kernel_size=3, padding=1) ] return nn.ModuleList(loc_layers), nn.ModuleList(conf_layers) def init_weights(self): pass def forward(self, feed_dict): img = feed_dict['img'] source_feats = self.feature_extractor(img) loc_preds = [] conf_preds = [] featmap_shapes = [] # apply multibox head to source layers for (x, l, c) in zip(source_feats, self.loc_layers, self.conf_layers): loc_preds.append(l(x).permute(0, 2, 3, 1).contiguous()) conf_preds.append(c(x).permute(0, 2, 3, 1).contiguous()) featmap_shapes.append(x.size()[-2:]) # import ipdb # ipdb.set_trace() loc_preds = torch.cat([o.view(o.size(0), -1) for o in loc_preds], 1) conf_preds = torch.cat([o.view(o.size(0), -1) for o in conf_preds], 1) probs = F.softmax( conf_preds.view(conf_preds.size(0), -1, self.n_classes), dim=-1) loc_preds = loc_preds.view(loc_preds.size(0), -1, 4) # import ipdb # ipdb.set_trace() anchors = self.anchor_generator.generate_pyramid(featmap_shapes) # anchors = self.priorsbox.forward(featmap_shapes) # import ipdb # ipdb.set_trace() rois_batch_inds = torch.zeros_like(loc_preds[:, :, -1:]) rois_batch = torch.cat([rois_batch_inds, anchors.unsqueeze(0)], dim=-1) second_rpn_anchors = anchors.unsqueeze(0) rcnn_3d = torch.zeros_like(loc_preds) prediction_dict = { 'rcnn_bbox_preds': loc_preds, 'rcnn_cls_scores': conf_preds, 'anchors': anchors, 'rcnn_cls_probs': probs, 'rois_batch': rois_batch, 'second_rpn_anchors': second_rpn_anchors, 'rcnn_3d': rcnn_3d } return prediction_dict def loss(self, prediction_dict, feed_dict): # import ipdb # ipdb.set_trace() # loss for cls loss_dict = {} gt_boxes = feed_dict['gt_boxes'] anchors = prediction_dict['anchors'] ################################# # target assigner ################################ # no need gt labels here,it just a binary classifcation problem # import ipdb # ipdb.set_trace() rpn_cls_targets, rpn_reg_targets, \ rpn_cls_weights, rpn_reg_weights = \ self.target_assigner.assign(anchors, gt_boxes, gt_labels=None) ################################ # subsample ################################ pos_indicator = rpn_reg_weights > 0 indicator = rpn_cls_weights > 0 rpn_cls_probs = prediction_dict['rcnn_cls_probs'][:, :, 1] cls_criterion = rpn_cls_probs batch_sampled_mask = self.sampler.subsample_batch( self.batch_size, pos_indicator, criterion=cls_criterion, indicator=indicator) batch_sampled_mask = batch_sampled_mask.type_as(rpn_cls_weights) rpn_cls_weights = rpn_cls_weights * batch_sampled_mask rpn_reg_weights = rpn_reg_weights * batch_sampled_mask num_cls_coeff = (rpn_cls_weights > 0).sum(dim=1) num_reg_coeff = (rpn_reg_weights > 0).sum(dim=1) # check # assert num_cls_coeff, 'bug happens' # assert num_reg_coeff, 'bug happens' if num_cls_coeff == 0: num_cls_coeff = torch.ones([]).type_as(num_cls_coeff) if num_reg_coeff == 0: num_reg_coeff = torch.ones([]).type_as(num_reg_coeff) # cls loss rpn_cls_score = prediction_dict['rcnn_cls_scores'] # rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score, rpn_cls_targets) rpn_cls_loss = self.conf_loss( rpn_cls_score.view(-1, 2), rpn_cls_targets.view(-1)) rpn_cls_loss = rpn_cls_loss.view_as(rpn_cls_weights) rpn_cls_loss *= rpn_cls_weights rpn_cls_loss = rpn_cls_loss.sum(dim=1) / num_cls_coeff.float() # bbox loss # shape(N,num,4) rpn_bbox_preds = prediction_dict['rcnn_bbox_preds'] # rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) # rpn_bbox_preds = rpn_bbox_preds.view(rpn_bbox_preds.shape[0], -1, 4) # import ipdb # ipdb.set_trace() rpn_reg_loss = self.loc_loss(rpn_bbox_preds, rpn_reg_targets) rpn_reg_loss *= rpn_reg_weights.unsqueeze(-1).expand(-1, -1, 4) rpn_reg_loss = rpn_reg_loss.view(rpn_reg_loss.shape[0], -1).sum( dim=1) / num_reg_coeff.float() prediction_dict['rcnn_reg_weights'] = rpn_reg_weights[ batch_sampled_mask > 0] loss_dict['rpn_cls_loss'] = rpn_cls_loss loss_dict['rpn_bbox_loss'] = rpn_reg_loss # recall final_boxes = self.bbox_coder.decode_batch(rpn_bbox_preds, anchors) self.target_assigner.assign(final_boxes, gt_boxes) return loss_dict
class RPNModel(Model): def init_param(self, model_config): self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] self.rpn_batch_size = model_config['rpn_batch_size'] self.use_focal_loss = model_config['use_focal_loss'] # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) # self.sampler = BalancedSampler(model_config['sampler_config']) self.sampler = DetectionSampler(model_config['sampler_config']) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = TargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder self.use_iou = model_config.get('use_iou') def init_weights(self): self.truncated = False Filler.normal_init(self.rpn_conv_cls, 0, 0.01, self.truncated) # Filler.normal_init(self.rpn_conv_bbox, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_cls_score, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_bbox_pred, 0, 0.01, self.truncated) def init_modules(self): # define the convrelu layers processing input feature map self.rpn_conv_cls = nn.Conv2d( self.in_channels, 512, 3, 1, 1, bias=True) # self.rpn_conv_bbox = nn.Conv2d( # self.in_channels, 512, 3, 1, 1, bias=True) # define bg/fg classifcation score layer self.rpn_cls_score = nn.Conv2d(512, self.nc_score_out, 1, 1, 0) # define anchor box offset prediction layer if self.use_score: bbox_feat_channels = 512 + 2 self.nc_bbox_out /= self.num_anchors else: bbox_feat_channels = 512 self.rpn_bbox_pred = nn.Conv2d(bbox_feat_channels, self.nc_bbox_out, 1, 1, 0) # bbox self.rpn_bbox_loss = nn.modules.loss.SmoothL1Loss(reduce=False) # cls if self.use_focal_loss: self.rpn_cls_loss = FocalLoss( 2, alpha=0.2, gamma=2, auto_alpha=False) else: self.rpn_cls_loss = functools.partial( F.cross_entropy, reduce=False) def generate_new_anchors(self, anchors): # import ipdb # ipdb.set_trace() anchor_size = 2 # anchors_w = anchors[:, :, 2] - anchors[:, :, 0] + 1 # anchors_h = anchors[:, :, 3] - anchors[:, :, 1] + 1 center_x = (anchors[:, 2] + anchors[:, 0]) / 2 center_y = (anchors[:, 3] + anchors[:, 1]) / 2 # new anchors has the same center as old ones min_x = center_x - (anchor_size - 1) / 2 min_y = center_y - (anchor_size - 1) / 2 max_x = center_x + (anchor_size - 1) / 2 max_y = center_y + (anchor_size - 1) / 2 return torch.stack([min_x, min_y, max_x, max_y], dim=-1) def generate_proposal(self, rpn_cls_probs, anchors, rpn_bbox_preds, im_info): # TODO create a new Function """ Args: rpn_cls_probs: FloatTensor,shape(N,2*num_anchors,H,W) rpn_bbox_preds: FloatTensor,shape(N,num_anchors*4,H,W) anchors: FloatTensor,shape(N,4,H,W) Returns: proposals_batch: FloatTensor, shape(N,post_nms_topN,4) fg_probs_batch: FloatTensor, shape(N,post_nms_topN) """ # assert len( # rpn_bbox_preds) == 1, 'just one feature maps is supported now' # rpn_bbox_preds = rpn_bbox_preds[0] # do not backward anchors = anchors rpn_cls_probs = rpn_cls_probs.detach() rpn_bbox_preds = rpn_bbox_preds.detach() batch_size = rpn_bbox_preds.shape[0] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(batch_size, -1, 4) # apply deltas to anchors to decode # loop here due to many features maps # proposals = [] # for rpn_bbox_preds_single_map, anchors_single_map in zip( # rpn_bbox_preds, anchors): # proposals.append( # self.bbox_coder.decode(rpn_bbox_preds_single_map, # anchors_single_map)) # proposals = torch.cat(proposals, dim=1) # make anchors small new_anchors = self.generate_new_anchors(anchors) proposals = self.bbox_coder.decode_batch(rpn_bbox_preds, new_anchors) # filer and clip proposals = box_ops.clip_boxes(proposals, im_info) # fg prob fg_probs = rpn_cls_probs[:, self.num_anchors:, :, :] fg_probs = fg_probs.permute(0, 2, 3, 1).contiguous().view(batch_size, -1) # sort fg _, fg_probs_order = torch.sort(fg_probs, dim=1, descending=True) # fg_probs_batch = torch.zeros(batch_size, # self.post_nms_topN).type_as(rpn_cls_probs) proposals_batch = torch.zeros(batch_size, self.post_nms_topN, 4).type_as(rpn_bbox_preds) proposals_order = torch.zeros( batch_size, self.post_nms_topN).fill_(-1).type_as(fg_probs_order) for i in range(batch_size): proposals_single = proposals[i] fg_probs_single = fg_probs[i] fg_order_single = fg_probs_order[i] # pre nms if self.pre_nms_topN > 0: fg_order_single = fg_order_single[:self.pre_nms_topN] proposals_single = proposals_single[fg_order_single] fg_probs_single = fg_probs_single[fg_order_single] # nms keep_idx_i = nms( torch.cat((proposals_single, fg_probs_single.unsqueeze(1)), 1), self.nms_thresh) keep_idx_i = keep_idx_i.long().view(-1) # post nms if self.post_nms_topN > 0: keep_idx_i = keep_idx_i[:self.post_nms_topN] proposals_single = proposals_single[keep_idx_i, :] fg_probs_single = fg_probs_single[keep_idx_i] fg_order_single = fg_order_single[keep_idx_i] # padding 0 at the end. num_proposal = keep_idx_i.numel() proposals_batch[i, :num_proposal, :] = proposals_single # fg_probs_batch[i, :num_proposal] = fg_probs_single proposals_order[i, :num_proposal] = fg_order_single return proposals_batch, proposals_order def forward(self, bottom_blobs): base_feat = bottom_blobs['base_feat'] batch_size = base_feat.shape[0] gt_boxes = bottom_blobs['gt_boxes'] im_info = bottom_blobs['im_info'] # separate cls featmap and bbox featmap # rpn conv rpn_conv_cls = F.relu(self.rpn_conv_cls(base_feat), inplace=True) # rpn cls score # shape(N,2*num_anchors,H,W) rpn_cls_scores = self.rpn_cls_score(rpn_conv_cls) # rpn cls prob shape(N,2*num_anchors,H,W) rpn_cls_score_reshape = rpn_cls_scores.view(batch_size, 2, -1) rpn_cls_probs = F.softmax(rpn_cls_score_reshape, dim=1) rpn_cls_probs = rpn_cls_probs.view_as(rpn_cls_scores) # import ipdb # ipdb.set_trace() # rpn bbox pred # shape(N,4*num_anchors,H,W) # if self.use_score: # # shape (N,2,num_anchoros*H*W) # rpn_cls_scores = rpn_cls_score_reshape.permute(0, 2, 1) # rpn_bbox_preds = [] # for i in range(self.num_anchors): # rpn_bbox_feat = torch.cat( # [rpn_conv, rpn_cls_scores[:, ::self.num_anchors, :, :]], # dim=1) # rpn_bbox_preds.append(self.rpn_bbox_pred(rpn_bbox_feat)) # rpn_bbox_preds = torch.cat(rpn_bbox_preds, dim=1) # else: # # get rpn offsets to the anchor boxes # rpn_bbox_preds = self.rpn_bbox_pred(rpn_conv) # # rpn_bbox_preds = [rpn_bbox_preds] # rpn_conv_bbox = F.relu(self.rpn_conv_bbox(base_feat), inplace=True) # shared with cls rpn_bbox_preds = self.rpn_bbox_pred(rpn_conv_cls) # generate anchors feature_map_list = [base_feat.size()[-2:]] anchors = self.anchor_generator.generate( feature_map_list, input_size=im_info[0][:-1]) ############################### # Proposal ############################### # note that proposals_order is used for track transform of propsoals proposals_batch, proposals_order = self.generate_proposal( rpn_cls_probs, anchors, rpn_bbox_preds, im_info) batch_idx = torch.arange(batch_size).view(batch_size, 1).expand( -1, proposals_batch.shape[1]).type_as(proposals_batch) rois_batch = torch.cat((batch_idx.unsqueeze(-1), proposals_batch), dim=2) rpn_cls_scores = rpn_cls_scores.view(batch_size, 2, -1, rpn_cls_scores.shape[2], rpn_cls_scores.shape[3]) rpn_cls_scores = rpn_cls_scores.permute( 0, 3, 4, 2, 1).contiguous().view(batch_size, -1, 2) # postprocess rpn_cls_probs = rpn_cls_probs.view( batch_size, 2, -1, rpn_cls_probs.shape[2], rpn_cls_probs.shape[3]) rpn_cls_probs = rpn_cls_probs.permute(0, 3, 4, 2, 1).contiguous().view( batch_size, -1, 2) predict_dict = { 'proposals_batch': proposals_batch, 'rpn_cls_scores': rpn_cls_scores, 'rois_batch': rois_batch, 'anchors': anchors, # used for loss 'rpn_bbox_preds': rpn_bbox_preds, 'rpn_cls_probs': rpn_cls_probs, 'proposals_order': proposals_order, } return predict_dict def loss(self, prediction_dict, feed_dict): # loss for cls loss_dict = {} gt_boxes = feed_dict['gt_boxes'] anchors = prediction_dict['anchors'] # small anchors new_anchors = self.generate_new_anchors(anchors) # assert len(anchors) == 1, 'just one feature maps is supported now' # anchors = anchors[0] ################################# # target assigner ################################ # no need gt labels here,it just a binary classifcation problem # import ipdb # ipdb.set_trace() rpn_cls_targets, rpn_reg_targets, \ rpn_cls_weights, rpn_reg_weights, stats = \ self.target_assigner.assign(anchors, gt_boxes, new_anchors, gt_labels=None) ################################ # subsample ################################ pos_indicator = rpn_reg_weights > 0 indicator = rpn_cls_weights > 0 rpn_cls_probs = prediction_dict['rpn_cls_probs'][:, :, 1] cls_criterion = rpn_cls_probs batch_sampled_mask = self.sampler.subsample_batch( self.rpn_batch_size, pos_indicator, criterion=cls_criterion, indicator=indicator) batch_sampled_mask = batch_sampled_mask.type_as(rpn_cls_weights) rpn_cls_weights = rpn_cls_weights * batch_sampled_mask rpn_reg_weights = rpn_reg_weights * batch_sampled_mask num_cls_coeff = (rpn_cls_weights > 0).sum(dim=1) num_reg_coeff = (rpn_reg_weights > 0).sum(dim=1) # check # assert num_cls_coeff, 'bug happens' # assert num_reg_coeff, 'bug happens' if num_cls_coeff == 0: num_cls_coeff = torch.ones([]).type_as(num_cls_coeff) if num_reg_coeff == 0: num_reg_coeff = torch.ones([]).type_as(num_reg_coeff) # cls loss rpn_cls_score = prediction_dict['rpn_cls_scores'] # rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score, rpn_cls_targets) rpn_cls_loss = self.rpn_cls_loss( rpn_cls_score.view(-1, 2), rpn_cls_targets.view(-1)) rpn_cls_loss = rpn_cls_loss.view_as(rpn_cls_weights) rpn_cls_loss *= rpn_cls_weights rpn_cls_loss = rpn_cls_loss.sum(dim=1) / num_cls_coeff.float() # bbox loss # shape(N,num,4) rpn_bbox_preds = prediction_dict['rpn_bbox_preds'] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(rpn_bbox_preds.shape[0], -1, 4) rpn_reg_loss = self.rpn_bbox_loss(rpn_bbox_preds, rpn_reg_targets) rpn_reg_loss *= rpn_reg_weights.unsqueeze(-1).expand(-1, -1, 4) rpn_reg_loss = rpn_reg_loss.view(rpn_reg_loss.shape[0], -1).sum( dim=1) / num_reg_coeff.float() loss_dict['rpn_cls_loss'] = rpn_cls_loss loss_dict['rpn_bbox_loss'] = rpn_reg_loss return loss_dict
class DistanceRPNModel(Model): def init_param(self, model_config): self.in_channels = model_config['din'] self.post_nms_topN = model_config['post_nms_topN'] self.pre_nms_topN = model_config['pre_nms_topN'] self.nms_thresh = model_config['nms_thresh'] self.use_score = model_config['use_score'] self.rpn_batch_size = model_config['rpn_batch_size'] self.use_focal_loss = model_config['use_focal_loss'] # sampler # self.sampler = HardNegativeSampler(model_config['sampler_config']) self.sampler = DetectionSampler(model_config['sampler_config']) # can not use hem here # self.sampler = BalancedSampler(model_config['sampler_config']) # anchor generator self.anchor_generator = AnchorGenerator( model_config['anchor_generator_config']) self.num_anchors = self.anchor_generator.num_anchors self.nc_bbox_out = 4 * self.num_anchors self.nc_score_out = self.num_anchors * 2 # target assigner self.target_assigner = DistanceTargetAssigner( model_config['target_assigner_config']) # bbox coder self.bbox_coder = self.target_assigner.bbox_coder def init_weights(self): self.truncated = False Filler.normal_init(self.rpn_conv, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_cls_score, 0, 0.01, self.truncated) Filler.normal_init(self.rpn_bbox_pred, 0, 0.01, self.truncated) def init_modules(self): # define the convrelu layers processing input feature map self.rpn_conv = nn.Conv2d(self.in_channels, 512, 3, 1, 1, bias=True) # define bg/fg classifcation score layer self.rpn_cls_score = nn.Conv2d(512, self.nc_score_out, 1, 1, 0) # define anchor box offset prediction layer if self.use_score: bbox_feat_channels = 512 + 2 self.nc_bbox_out /= self.num_anchors else: bbox_feat_channels = 512 self.rpn_bbox_pred = nn.Conv2d(bbox_feat_channels, self.nc_bbox_out, 1, 1, 0) # bbox self.rpn_bbox_loss = nn.modules.loss.SmoothL1Loss(reduce=False) # cls if self.use_focal_loss: self.rpn_cls_loss = FocalLoss(2) else: self.rpn_cls_loss = functools.partial( F.cross_entropy, reduce=False) # self.rpn_cls_loss = nn.MSELoss(reduce=False) # self.distance_similarity_calc = DistanceSimilarityCalc() def generate_proposal(self, rpn_cls_probs, anchors, rpn_bbox_preds, im_info): # TODO create a new Function """ Args: rpn_cls_probs: FloatTensor,shape(N,2*num_anchors,H,W) rpn_bbox_preds: FloatTensor,shape(N,num_anchors*4,H,W) anchors: FloatTensor,shape(N,4,H,W) Returns: proposals_batch: FloatTensor, shape(N,post_nms_topN,4) fg_probs_batch: FloatTensor, shape(N,post_nms_topN) """ # assert len( # rpn_bbox_preds) == 1, 'just one feature maps is supported now' # rpn_bbox_preds = rpn_bbox_preds[0] anchors = anchors[0] # do not backward anchors = anchors rpn_cls_probs = rpn_cls_probs.detach() rpn_bbox_preds = rpn_bbox_preds.detach() batch_size = rpn_bbox_preds.shape[0] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(batch_size, -1, 4) # fg prob gate = rpn_cls_probs[:, self.num_anchors:, :, :] gate = gate.permute(0, 2, 3, 1).contiguous().view(batch_size, -1) fg_probs, distance = self.get_rpn_cls_probs( rpn_bbox_preds, anchors=None) fg_probs[gate < 0.5] = 0 distance[gate < 0.5] = 1e5 # apply deltas to anchors to decode # loop here due to many features maps # proposals = [] # for rpn_bbox_preds_single_map, anchors_single_map in zip( # rpn_bbox_preds, anchors): # proposals.append( # self.bbox_coder.decode(rpn_bbox_preds_single_map, # anchors_single_map)) # proposals = torch.cat(proposals, dim=1) proposals = self.bbox_coder.decode_batch(rpn_bbox_preds, anchors) # filer and clip proposals = box_ops.clip_boxes(proposals, im_info) # fg prob # sort fg _, fg_probs_order = torch.sort(fg_probs, dim=1, descending=True) # fg_probs_batch = torch.zeros(batch_size, # self.post_nms_topN).type_as(rpn_cls_probs) proposals_batch = torch.zeros(batch_size, self.post_nms_topN, 4).type_as(rpn_bbox_preds) proposals_order = torch.zeros( batch_size, self.post_nms_topN).fill_(-1).type_as(fg_probs_order) for i in range(batch_size): proposals_single = proposals[i] fg_probs_single = fg_probs[i] fg_order_single = fg_probs_order[i] # pre nms if self.pre_nms_topN > 0: fg_order_single = fg_order_single[:self.pre_nms_topN] proposals_single = proposals_single[fg_order_single] fg_probs_single = fg_probs_single[fg_order_single] # nms keep_idx_i = nms( torch.cat((proposals_single, fg_probs_single.unsqueeze(1)), 1), self.nms_thresh) keep_idx_i = keep_idx_i.long().view(-1) # post nms if self.post_nms_topN > 0: keep_idx_i = keep_idx_i[:self.post_nms_topN] proposals_single = proposals_single[keep_idx_i, :] fg_probs_single = fg_probs_single[keep_idx_i] fg_order_single = fg_order_single[keep_idx_i] # padding 0 at the end. num_proposal = keep_idx_i.numel() proposals_batch[i, :num_proposal, :] = proposals_single # fg_probs_batch[i, :num_proposal] = fg_probs_single proposals_order[i, :num_proposal] = fg_order_single # row = torch.arange(0, batch_size).type_as(proposals_order) # fg_probs = fg_probs[row, proposals_order.view(-1)].view_as( # proposals_order) return proposals_batch, proposals_order, fg_probs, distance def get_rpn_cls_probs(self, bbox_pred, anchors=None): """ Note that all inputs have no gradients Args: bbox_pred: shape (N,M,4) anchors: shape (M,4) Returns: distance: shape(N,M) """ # shape(N,M,4) # distances = self.distance_similarity_calc.compare_batch(bbox, gt_boxes) # anchors = anchors.expand_as(bbox_pred) # widths = anchors[:, :, 2] - anchors[:, :, 0] + 1.0 # heights = anchors[:, :, 3] - anchors[:, :, 1] + 1.0 # dx = bbox_pred[:, :, 0] * widths # dy = bbox_pred[:, :, 1] * heights dx = bbox_pred[:, :, 0] dy = bbox_pred[:, :, 1] distance = torch.sqrt(dx * dx + dy * dy) theta = 1e-5 return 1.0 / (distance + theta), distance def forward(self, bottom_blobs): base_feat = bottom_blobs['base_feat'] batch_size = base_feat.shape[0] gt_boxes = bottom_blobs['gt_boxes'] im_info = bottom_blobs['im_info'] # rpn conv rpn_conv = F.relu(self.rpn_conv(base_feat), inplace=True) # rpn bbox pred # shape(N,4*num_anchors,H,W) rpn_bbox_preds = self.rpn_bbox_pred(rpn_conv) # generate anchors feature_map_list = [base_feat.size()[-2:]] anchors = self.anchor_generator.generate(feature_map_list) rpn_cls_scores = self.rpn_cls_score(rpn_conv) # softmax rpn_cls_score_reshape = rpn_cls_scores.view(batch_size, 2, -1) rpn_cls_probs = F.softmax(rpn_cls_score_reshape, dim=1) rpn_cls_probs = rpn_cls_probs.view_as(rpn_cls_scores) # use distance to instead of rpn_cls_probs # rpn_cls_probs = self.get_rpn_cls_probs(rpn_bbox_preds) ############################### # Proposal ############################### # note that proposals_order is used for track transform of propsoals proposals_batch, proposals_order, fg_probs, distance = self.generate_proposal( rpn_cls_probs, anchors, rpn_bbox_preds, im_info) batch_idx = torch.arange(batch_size).view(batch_size, 1).expand( -1, proposals_batch.shape[1]).type_as(proposals_batch) rois_batch = torch.cat((batch_idx.unsqueeze(-1), proposals_batch), dim=2) if self.training: rois_batch = self.append_gt(rois_batch, gt_boxes) # postprocess rpn_cls_scores = rpn_cls_scores.view(batch_size, 2, -1, rpn_cls_scores.shape[2], rpn_cls_scores.shape[3]) rpn_cls_scores = rpn_cls_scores.permute( 0, 3, 4, 2, 1).contiguous().view(batch_size, -1, 2) predict_dict = { 'proposals_batch': proposals_batch, 'rois_batch': rois_batch, 'anchors': anchors, # used for loss 'rpn_bbox_preds': rpn_bbox_preds, 'rpn_cls_scores': rpn_cls_scores, 'proposals_order': proposals_order, 'fg_probs': fg_probs, 'distance': distance } return predict_dict def append_gt(self, rois_batch, gt_boxes): ################################ # append gt_boxes to rois_batch for losses ################################ # may be some bugs here gt_boxes_append = torch.zeros(gt_boxes.shape[0], gt_boxes.shape[1], 5).type_as(gt_boxes) gt_boxes_append[:, :, 1:5] = gt_boxes[:, :, :4] # cat gt_boxes to rois_batch rois_batch = torch.cat([rois_batch, gt_boxes_append], dim=1) return rois_batch def loss(self, prediction_dict, feed_dict): # loss for cls loss_dict = {} gt_boxes = feed_dict['gt_boxes'] anchors = prediction_dict['anchors'] assert len(anchors) == 1, 'just one feature maps is supported now' anchors = anchors[0] ################################# # target assigner ################################ # no need gt labels here,it just a binary classifcation problem # import ipdb # ipdb.set_trace() rpn_cls_targets, rpn_reg_targets, \ rpn_cls_weights, rpn_reg_weights = \ self.target_assigner.assign(anchors, gt_boxes, gt_labels=None) ################################ # subsample ################################ rpn_cls_probs = prediction_dict['fg_probs'] pos_indicator = rpn_cls_targets > 0 indicator = rpn_cls_weights > 0 cls_criterion = rpn_cls_probs batch_sampled_mask = self.sampler.subsample_batch( self.rpn_batch_size, pos_indicator, criterion=cls_criterion, indicator=indicator) batch_sampled_mask = batch_sampled_mask.type_as(rpn_cls_weights) rpn_cls_weights = rpn_cls_weights * batch_sampled_mask rpn_reg_weights = rpn_reg_weights * batch_sampled_mask num_cls_coeff = (rpn_cls_weights > 0).sum(dim=1) num_reg_coeff = (rpn_reg_weights > 0).sum(dim=1) # check # assert num_cls_coeff, 'bug happens' # assert num_reg_coeff, 'bug happens' if num_cls_coeff == 0: num_cls_coeff = torch.ones([]).type_as(num_cls_coeff) if num_reg_coeff == 0: num_reg_coeff = torch.ones([]).type_as(num_reg_coeff) # cls loss rpn_cls_scores = prediction_dict['rpn_cls_scores'] rpn_cls_loss = self.rpn_cls_loss(rpn_cls_scores, rpn_cls_targets) rpn_cls_loss = rpn_cls_loss.view_as(rpn_cls_weights) rpn_cls_loss *= rpn_cls_weights rpn_cls_loss = rpn_cls_loss.sum(dim=1) / num_cls_coeff.float() # rpn_cls_probs = prediction_dict['rpn_cls_probs'] # fg_rpn_cls_probs = rpn_cls_probs.view(-1, 2)[:, 1] # # exp # fg_rpn_cls_probs = torch.exp(fg_rpn_cls_probs) # rpn_cls_targets = torch.exp(rpn_cls_targets) # # rpn_cls_loss = self.rpn_cls_loss(rpn_cls_score, rpn_cls_targets) # rpn_cls_loss = self.rpn_cls_loss(fg_rpn_cls_probs, # rpn_cls_targets.view(-1)) # rpn_cls_loss = rpn_cls_loss.view_as(rpn_cls_weights) # rpn_cls_loss *= rpn_cls_weights # rpn_cls_loss = rpn_cls_loss.sum(dim=1) / num_cls_coeff.float() # bbox loss # shape(N,num,4) rpn_bbox_preds = prediction_dict['rpn_bbox_preds'] rpn_bbox_preds = rpn_bbox_preds.permute(0, 2, 3, 1).contiguous() # shape(N,H*W*num_anchors,4) rpn_bbox_preds = rpn_bbox_preds.view(rpn_bbox_preds.shape[0], -1, 4) rpn_reg_loss = self.rpn_bbox_loss(rpn_bbox_preds, rpn_reg_targets) rpn_reg_loss *= rpn_reg_weights.unsqueeze(-1).expand(-1, -1, 4) rpn_reg_loss = rpn_reg_loss.view(rpn_reg_loss.shape[0], -1).sum( dim=1) / num_reg_coeff.float() loss_dict['rpn_cls_loss'] = rpn_cls_loss loss_dict['rpn_bbox_loss'] = rpn_reg_loss return loss_dict