def _single_gpu_build_func(model): """Builds the model on a single GPU. Can be called in a loop over GPUs with name and device scoping to create a data parallel model.""" if model.train: train=model.train switch_to_teacher() model.train=False init_params=model.init_params with c2_utils.NamedTeacherScope(): blobs, dim, spatial_scales = get_func(cfg.MODEL.CONV_BODY)(model) retinanet_heads.add_fpn_retinanet_outputs( model, blobs, dim, spatial_scales ) model.train=train model.init_params=init_params switch_to_student() blobs, dim, spatial_scales = add_conv_body_func(model) if not model.train: model.conv_body_net = model.net.Clone('conv_body_net') retinanet_heads.add_fpn_retinanet_outputs( model, blobs, dim, spatial_scales ) if model.train: loss_gradients = retinanet_heads.add_fpn_retinanet_losses( model ) loss_gradients_distill=retinanet_heads.add_distill_loss(model,'','teacher/') loss_gradients.update(loss_gradients_distill) return loss_gradients if model.train else None
def _single_gpu_build_func(model): """Builds the model on a single GPU. Can be called in a loop over GPUs with name and device scoping to create a data parallel model.""" blobs, dim, spatial_scales = add_conv_body_func(model) retinanet_heads.add_fpn_retinanet_outputs(model, blobs, dim, spatial_scales) if model.train: loss_gradients = retinanet_heads.add_fpn_retinanet_losses(model) return loss_gradients if model.train else None
def _single_gpu_build_func(model): """Builds the model on a single GPU. Can be called in a loop over GPUs with name and device scoping to create a data parallel model.""" blobs, dim, spatial_scales = add_conv_body_func(model) retinanet_heads.add_fpn_retinanet_outputs( model, blobs, dim, spatial_scales ) if model.train: loss_gradients = retinanet_heads.add_fpn_retinanet_losses( model ) return loss_gradients if model.train else None
def _forward(self, data, im_info, roidb=None, **rpn_kwargs): im_data = data if self.training: roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb)) device_id = im_data.get_device() return_dict = {} # A dict to collect return variables blob_conv = self.Conv_Body(im_data) # if self.training: # # can be used to infer fg/bg ratio # return_dict['rois_label'] = rpn_ret['labels_int32'] if cfg.RPN.RPN_ON: rpn_ret = self.RPN(blob_conv, im_info, roidb) if cfg.FPN.FPN_ON and cfg.FAST_RCNN.ROI_BOX_HEAD is not '': # Retain only the blobs that will be used for RoI heads. `blob_conv` may include # extra blobs that are used for RPN proposals, but not for RoI # heads. blob_conv = blob_conv[-self.num_roi_levels:] if not self.training: return_dict['blob_conv'] = blob_conv if not cfg.MODEL.RPN_ONLY: if cfg.FAST_RCNN.ROI_BOX_HEAD is not '': if cfg.MODEL.SHARE_RES5 and self.training: box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret) else: box_feat = self.Box_Head(blob_conv, rpn_ret) else: box_feat = blob_conv cls_score, bbox_pred = self.Box_Outs(box_feat) else: # TODO: complete the returns for RPN only situation pass if self.training: return_dict['losses'] = {} return_dict['metrics'] = {} if cfg.FAST_RCNN.ROI_BOX_HEAD is not '': # rpn loss rpn_kwargs.update(dict((k, rpn_ret[k]) for k in rpn_ret.keys() if ( k.startswith('rpn_cls_logits') or k.startswith('rpn_bbox_pred')))) loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses( **rpn_kwargs) if cfg.FPN.FPN_ON: for i, lvl in enumerate( range(cfg.FPN.RPN_MIN_LEVEL, cfg.FPN.RPN_MAX_LEVEL + 1)): return_dict['losses']['loss_rpn_cls_fpn%d' % lvl] = loss_rpn_cls[i] return_dict['losses']['loss_rpn_bbox_fpn%d' % lvl] = loss_rpn_bbox[i] else: return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox # bbox loss loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses( cls_score, bbox_pred, rpn_ret['labels_int32'], rpn_ret['bbox_targets'], rpn_ret['bbox_inside_weights'], rpn_ret['bbox_outside_weights']) return_dict['losses']['loss_cls'] = loss_cls return_dict['losses']['loss_bbox'] = loss_bbox return_dict['metrics']['accuracy_cls'] = accuracy_cls if cfg.RETINANET.RETINANET_ON: loss_retnet_cls, loss_retnet_bbox = retinanet_heads.add_fpn_retinanet_losses( cls_score, bbox_pred, **rpn_kwargs) for i, lvl in enumerate( range(cfg.FPN.RPN_MIN_LEVEL, cfg.FPN.RPN_MAX_LEVEL + 1)): return_dict['losses']['loss_retnet_cls_fpn%d' % lvl] = loss_retnet_cls[i] return_dict['losses']['loss_retnet_bbox_fpn%d' % lvl] = loss_retnet_bbox[i] if cfg.MODEL.MASK_ON: if getattr(self.Mask_Head, 'SHARE_RES5', False): mask_feat = self.Mask_Head( res5_feat, rpn_ret, roi_has_mask_int32=rpn_ret['roi_has_mask_int32']) else: mask_feat = self.Mask_Head(blob_conv, rpn_ret) mask_pred = self.Mask_Outs(mask_feat) # return_dict['mask_pred'] = mask_pred # mask loss loss_mask = mask_rcnn_heads.mask_rcnn_losses( mask_pred, rpn_ret['masks_int32']) return_dict['losses']['loss_mask'] = loss_mask if cfg.MODEL.KEYPOINTS_ON: if getattr(self.Keypoint_Head, 'SHARE_RES5', False): # No corresponding keypoint head implemented yet (Neither in Detectron) # Also, rpn need to generate the label # 'roi_has_keypoints_int32' kps_feat = self.Keypoint_Head( res5_feat, rpn_ret, roi_has_keypoints_int32=rpn_ret['roi_has_keypoint_int32']) else: kps_feat = self.Keypoint_Head(blob_conv, rpn_ret) kps_pred = self.Keypoint_Outs(kps_feat) # return_dict['keypoints_pred'] = kps_pred # keypoints loss if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS: loss_keypoints = keypoint_rcnn_heads.keypoint_losses( kps_pred, rpn_ret['keypoint_locations_int32'], rpn_ret['keypoint_weights']) else: loss_keypoints = keypoint_rcnn_heads.keypoint_losses( kps_pred, rpn_ret['keypoint_locations_int32'], rpn_ret['keypoint_weights'], rpn_ret['keypoint_loss_normalizer']) return_dict['losses']['loss_kps'] = loss_keypoints # pytorch0.4 bug on gathering scalar(0-dim) tensors for k, v in return_dict['losses'].items(): return_dict['losses'][k] = v.unsqueeze(0) for k, v in return_dict['metrics'].items(): return_dict['metrics'][k] = v.unsqueeze(0) else: # Testing if cfg.FAST_RCNN.ROI_BOX_HEAD is not '': return_dict['rois'] = rpn_ret['rois'] return_dict['cls_score'] = cls_score return_dict['bbox_pred'] = bbox_pred return return_dict