def build_generic_retinanet_model(model, add_conv_body_func, freeze_conv_body=False): # TODO(rbg): fold this function into build_generic_detection_model 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) 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 cfg.WSL.WSL_ON: retinanet_heads.add_wsl_fpn_retinanet_outputs( model, blobs, dim, spatial_scales) if model.train: loss_gradients = retinanet_heads.add_wsl_fpn_retinanet_losses( model) return loss_gradients if model.train else None if model.train: loss_gradients = retinanet_heads.add_fpn_retinanet_losses(model) return loss_gradients if model.train else None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_detection_model( model, add_conv_body_func, add_roi_box_head_func=None, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False ): def _single_gpu_build_func(model): """Build 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. """ # Add the conv body (called "backbone architecture" in papers) # E.g., ResNet-50, ResNet-50-FPN, ResNeXt-101-FPN, etc. blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, } if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv ) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv ) if not cfg.MODEL.RPN_ONLY: # Add the Fast R-CNN head head_loss_gradients['box'] = _add_fast_rcnn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv ) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_rfcn_model(model, add_conv_body_func, add_roi_box_head_func=None, dim_reduce=None, freeze_conv_body=False): # TODO(rbg): fold this function into build_generic_detection_model 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.""" blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, 'classification': None } if freeze_conv_body: for b in c2_utils.BlobReferenceList(blob_conv): model.StopGradient(b, b) if not model.train: model.conv_body_net = model.net.Clone('conv_body_net') if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv ) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv ) head_loss_gradients['box'] = _add_rfcn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv, dim_reduce ) #rfcn_heads.add_rfcn_outputs(model, blob_conv, dim_conv, dim_reduce, spatial_scale_conv) #if model.train: #head_loss_gradients['box'] = fast_rcnn_heads.add_fast_rcnn_losses(model) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None #loss_gradients = fast_rcnn_heads.add_fast_rcnn_losses(model) #return loss_gradients if model.train else None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_rfcn_model(model, add_conv_body_func, dim_reduce=None): # TODO(rbg): fold this function into build_generic_detection_model 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.""" blob, dim, spatial_scale = add_conv_body_func(model) if not model.train: model.conv_body_net = model.net.Clone('conv_body_net') rfcn_heads.add_rfcn_outputs(model, blob, dim, dim_reduce, spatial_scale) if model.train: loss_gradients = fast_rcnn_heads.add_fast_rcnn_losses(model) return loss_gradients if model.train else None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_rfcn_model(model, add_conv_body_func, dim_reduce=None): # TODO(rbg): fold this function into build_generic_detection_model 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.""" blob, dim, spatial_scale = add_conv_body_func(model) if not model.train: model.conv_body_net = model.net.Clone('conv_body_net') rfcn_heads.add_rfcn_outputs(model, blob, dim, dim_reduce, spatial_scale) if model.train: loss_gradients = fast_rcnn_heads.add_fast_rcnn_losses(model) return loss_gradients if model.train else None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_detection_model(model, add_conv_body_func, add_roi_box_head_func=None, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False): def _single_gpu_build_func(model): """Build 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. """ # Add the conv body (called "backbone architecture" in papers) # E.g., ResNet-50, ResNet-50-FPN, ResNeXt-101-FPN, etc. blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) if freeze_conv_body: for b in c2_utils.BlobReferenceList(blob_conv): model.StopGradient(b, b) if not model.train: # == inference # Create a net that can be used to execute the conv body on an image # (without also executing RPN or any other network heads) model.conv_body_net = model.net.Clone('conv_body_net') head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, 'box_pose': None, 'res_seg_human': None } if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv) # mask restop head_loss_gradients[ 'res_seg_human'] = mask_res_top_attention.add_mask_res_branch( model, blob_conv, 256) if not cfg.MODEL.RPN_ONLY: # Add the Fast R-CNN head head_loss_gradients['box'] = _add_fast_rcnn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv) # add FPN pose branch for shoes,legs # head_loss_gradients['box_pose'] = add_pose_fpn.add_pose_fpn_rcnn_head(model, blob_conv[:2], dim_conv, spatial_scale_conv) if cfg.MODEL.MASK_ON: # Add the mask head head_loss_gradients['mask'] = _add_roi_mask_head( model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv) if cfg.MODEL.KEYPOINTS_ON: # Add the keypoint head head_loss_gradients['keypoint'] = _add_roi_keypoint_head( model, add_roi_keypoint_head_func, blob_conv, dim_conv, spatial_scale_conv) # head_loss_gradients['res_seg_human'] = mask_res_top.add_mask_res_branch(model, blob_conv, 256) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_detection_model_for_feature_extraction( model, add_conv_body_func, add_roi_frcn_head_func, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False): 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.""" # For training we define one net that contains all ops # For inference, we split the graph into two nets: a standard fast r-cnn # net and a mask prediction net; the mask net is only applied to a # subset of high-scoring detections is_inference = not model.train head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, } # Add the conv body blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) if freeze_conv_body: for b in blob_ref_to_list(blob_conv): model.StopGradient(b, b) backbone_net = copy.deepcopy(model.net.Proto()) model.backbone = model.net.Clone('backbone') if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = model_builder._narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv) assert not cfg.MODEL.RPN_ONLY, 'only mask rcnn end2end implemented and tested' # Add the Fast R-CNN head head_loss_gradients[ 'box'], blob_cls, blob_bbox = _add_fast_rcnn_head_for_feature_extraction( model, add_roi_frcn_head_func, blob_conv, dim_conv, spatial_scale_conv) # Extract the head of faster rcnn net, store it as its own network, # then create a new network faster_rcnn_head model.faster_rnn_head, [blob_cls, blob_bbox] = SuffixNet('faster_rcnn_head', model.net, len(backbone_net.op), [blob_cls, blob_bbox]) if cfg.MODEL.MASK_ON: head_loss_gradients['mask'] = model_builder._add_roi_mask_head( model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv) # Add the keypoint branch assert not cfg.MODEL.KEYPOINTS_ON, 'keypoints not implemented, otherwise need to extend here' assert not model.train, 'not implemented, should add losses here' optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def create_model(train, cfg, output_dir): logger = logging.getLogger(__name__) start_iter = 0 checkpoints = {} weights_file = None if cfg.TRAIN.AUTO_RESUME: # Check for the final model (indicates training already finished) final_path = os.path.join(output_dir, 'model_final.pkl') if os.path.exists(final_path): logger.info('model_final.pkl exists; no need to train!') return None, None, None, {'final': final_path}, output_dir # Find the most recent checkpoint (highest iteration number) files = os.listdir(output_dir) for f in files: iter_string = re.findall(r'(?<=model_iter)\d+(?=\.pkl)', f) if len(iter_string) > 0: checkpoint_iter = int(iter_string[0]) if checkpoint_iter > start_iter: # Start one iteration immediately after the checkpoint iter start_iter = checkpoint_iter + 1 resume_weights_file = f if start_iter > 0: # Override the initialization weights with the found checkpoint weights_file = os.path.join(output_dir, resume_weights_file) logger.info( '========> Resuming from checkpoint {} at start iter {}'. format(weights_file, start_iter)) logger.info('Building model: {}'.format(cfg.MODEL.TYPE)) model = DetectionModelHelper(name="fusion", train=train, num_classes=cfg.MODEL.NUM_CLASSES, init_params=train) model.only_build_forward_pass = False model.target_gpu_id = gpu_id def _single_gpu_build_func(model): loss = {"loss": None} p = model.Conv("data_stage2", 'conv1_stage2', 40, 20, kernel=3, pad=1, stride=1, no_bias=1, weight_init=(cfg.MRCNN.CONV_INIT, { 'std': 0.001 })) p = model.AffineChannel(p, 'conv1_bn_stage2', dim=20, inplace=True) p = model.Relu(p, p) human_fc = model.Conv(p, 'conv2_stage2', 20, 20, kernel=1, pad=0, stride=1, weight_init=(cfg.MRCNN.CONV_INIT, { 'std': 0.001 })) if not model.train: model.net.NCHW2NHWC(human_fc, 'seg_score_NHWC_stage2') model.net.Softmax('seg_score_NHWC_stage2', 'probs_human_NHWC_stage2', axis=3) model.net.NHWC2NCHW('probs_human_NHWC_stage2', 'probs_human_NCHW_stage2') loss_gradient = None if model.train: model.net.NCHW2NHWC(human_fc, 'seg_score_NHWC_stage2') model.Reshape( 'seg_score_NHWC_stage2', ['seg_score_reshape_stage2', 'seg_score_old_shape_stage2'], shape=[-1, model.num_classes]) model.Reshape('gt_label_stage2', ['gt_label_reshape_stage2', 'gt_label_shape_stage2'], shape=[ -1, ]) probs_human, loss_human = model.net.SoftmaxWithLoss( ['seg_score_reshape_stage2', 'gt_label_reshape_stage2'], ['probs_human_stage2', 'loss_human_stage2'], scale=1. / cfg.NUM_GPUS) loss_gradient = blob_utils.get_loss_gradients(model, [loss_human]) model.AddLosses('loss_human_stage2') loss['loss'] = loss_gradient if model.train: loss_gradients = {} for lg in loss.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None optim.build_data_parallel_model(model, _single_gpu_build_func) # Performs random weight initialization as defined by the model workspace.RunNetOnce(model.param_init_net) return model, weights_file, start_iter, checkpoints
def build_generic_detection_model( model, add_conv_body_func, add_roi_box_head_func=None, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False ): def _single_gpu_build_func(model): """Build 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. """ # Add the conv body (called "backbone architecture" in papers) # E.g., ResNet-50, ResNet-50-FPN, ResNeXt-101-FPN, etc. blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) if freeze_conv_body: for b in c2_utils.BlobReferenceList(blob_conv): model.StopGradient(b, b) if not model.train: # == inference # Create a net that can be used to execute the conv body on an image # (without also executing RPN or any other network heads) model.conv_body_net = model.net.Clone('conv_body_net') head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, } if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv ) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv ) if not cfg.MODEL.RPN_ONLY: # Add the Fast R-CNN head head_loss_gradients['box'] = _add_fast_rcnn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv ) if cfg.MODEL.MASK_ON: # Add the mask head head_loss_gradients['mask'] = _add_roi_mask_head( model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv ) if cfg.MODEL.KEYPOINTS_ON: # Add the keypoint head head_loss_gradients['keypoint'] = _add_roi_keypoint_head( model, add_roi_keypoint_head_func, blob_conv, dim_conv, spatial_scale_conv ) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_detection_model(model, add_conv_body_func, add_roi_box_head_func=None, add_roi_cascade_head_func=None, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False): def _single_gpu_build_func(model): """Build 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. """ # Add the conv body (called "backbone architecture" in papers) # E.g., ResNet-50, ResNet-50-FPN, ResNeXt-101-FPN, etc. blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) if freeze_conv_body: for b in c2_utils.BlobReferenceList(blob_conv): model.StopGradient(b, b) if not model.train: # == inference # Create a net that can be used to execute the conv body on an image # (without also executing RPN or any other network heads) model.conv_body_net = model.net.Clone('conv_body_net') head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, } if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv) if not cfg.MODEL.RPN_ONLY: # Add the Fast R-CNN head head_loss_gradients['box'] = _add_fast_rcnn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv) if cfg.MODEL.CASCADE_ON: # Add the Cascade R-CNN head num_stage = cfg.CASCADE_RCNN.NUM_STAGE _check_for_cascade_rcnn() for stage in range(2, num_stage + 1): stage_name = '_{}'.format(stage) head_loss_gradients['box' + stage_name] = _add_cascade_rcnn_head( model, add_roi_cascade_head_func, blob_conv, dim_conv, spatial_scale_conv, stage) if cfg.MODEL.MASK_ON: # Add the mask head head_loss_gradients['mask'] = _add_roi_mask_head( model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv) if cfg.MODEL.KEYPOINTS_ON: # Add the keypoint head head_loss_gradients['keypoint'] = _add_roi_keypoint_head( model, add_roi_keypoint_head_func, blob_conv, dim_conv, spatial_scale_conv) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_detection_model( model, add_conv_body_func, add_roi_box_head_func=None, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False ): def _single_gpu_build_func(model): """Build 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. """ # Add the conv body (called "backbone architecture" in papers) # E.g., ResNet-50, ResNet-50-FPN, ResNeXt-101-FPN, etc. blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) if freeze_conv_body: for b in c2_utils.BlobReferenceList(blob_conv): model.StopGradient(b, b) if not model.train: # == inference # Create a net that can be used to execute the conv body on an image # (without also executing RPN or any other network heads) model.conv_body_net = model.net.Clone('conv_body_net') head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, 'image': None, 'instance': None, 'consistency': None, } if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv ) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv ) if not cfg.MODEL.RPN_ONLY: # Add the Fast R-CNN head head_loss_gradients['box'] = _add_fast_rcnn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv ) if cfg.MODEL.MASK_ON: # Add the mask head head_loss_gradients['mask'] = _add_roi_mask_head( model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv ) if cfg.MODEL.KEYPOINTS_ON: # Add the keypoint head head_loss_gradients['keypoint'] = _add_roi_keypoint_head( model, add_roi_keypoint_head_func, blob_conv, dim_conv, spatial_scale_conv ) if cfg.TRAIN.DOMAIN_ADAPTATION: # Add Image-level loss head_loss_gradients['image'] = _add_image_level_classifier(model, blob_conv, dim_conv, spatial_scale_conv) # Add Instance-level loss head_loss_gradients['instance'], blob_feats_rois_all, dim_feats_rois_all = _add_instance_level_classifier(model, blob_conv, dim_conv, spatial_scale_conv) # Add consistency regularization head_loss_gradients['consistency'] = _add_consistency_loss(model, blob_conv, dim_conv, blob_feats_rois_all, dim_feats_rois_all) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_detection_model(model, add_conv_body_func, add_roi_box_head_func=None, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False): def _single_gpu_build_func(model): """Build 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. 在一个gpu上创建模型 """ # Add the conv body (called "backbone architecture" in papers) # E.g., ResNet-50, ResNet-50-FPN, ResNeXt-101-FPN, etc. # 添加骨干网络 如ResNet.add_ResNet50_conv4_body # 对于fpn faster rcnn 101 # [sum5_subsmapled, sum5, ..., sum2] 256 [1 / 64, ..., 1 / 4] blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) # 冻结conv body,对每一个fpn层停止反向传播 if freeze_conv_body: for b in c2_utils.BlobReferenceList(blob_conv): model.StopGradient(b, b) if not model.train: # == inference # Create a net that can be used to execute the conv body on an image # (without also executing RPN or any other network heads) model.conv_body_net = model.net.Clone('conv_body_net') head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, } if cfg.RPN.RPN_ON: # Add the RPN head # 添加loss和grad head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv) if not cfg.MODEL.RPN_ONLY: # Add the Fast R-CNN head # 添加class和box regress损失 head_loss_gradients['box'] = _add_fast_rcnn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv) if cfg.MODEL.MASK_ON: # Add the mask head # 添加mask头 head_loss_gradients['mask'] = _add_roi_mask_head( model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv) if cfg.MODEL.KEYPOINTS_ON: # Add the keypoint head head_loss_gradients['keypoint'] = _add_roi_keypoint_head( model, add_roi_keypoint_head_func, blob_conv, dim_conv, spatial_scale_conv) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None optim.build_data_parallel_model(model, _single_gpu_build_func) return model
def build_generic_detection_model( model, add_conv_body_func, add_roi_box_head_func=None, add_roi_mask_head_func=None, add_roi_keypoint_head_func=None, freeze_conv_body=False ): def _single_gpu_build_func(model): """Build 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. """ # Add the conv body (called "backbone architecture" in papers) # E.g., ResNet-50, ResNet-50-FPN, ResNeXt-101-FPN, etc. blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model) if freeze_conv_body: for b in c2_utils.BlobReferenceList(blob_conv): model.StopGradient(b, b) if not model.train: # == inference # Create a net that can be used to execute the conv body on an image # (without also executing RPN or any other network heads) model.conv_body_net = model.net.Clone('conv_body_net') head_loss_gradients = { 'rpn': None, 'box': None, 'mask': None, 'keypoints': None, } if cfg.RPN.RPN_ON: # Add the RPN head head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs( model, blob_conv, dim_conv, spatial_scale_conv ) if cfg.FPN.FPN_ON: # After adding the RPN head, restrict FPN blobs and scales to # those used in the RoI heads blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels( blob_conv, spatial_scale_conv ) if not cfg.MODEL.RPN_ONLY: # Add the Fast R-CNN head head_loss_gradients['box'] = _add_fast_rcnn_head( model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv ) else: if cfg.MODEL.RPN_RA_AND_BP: roi_size = cfg.FAST_RCNN.ROI_XFORM_RESOLUTION model.RoIFeatureTransform( blob_conv, 'roi_feat', blob_rois='rois', method=cfg.FAST_RCNN.ROI_XFORM_METHOD, resolution=roi_size, sampling_ratio=cfg.FAST_RCNN.ROI_XFORM_SAMPLING_RATIO, spatial_scale=spatial_scale_conv ) if cfg.MODEL.MASK_ON: # Add the mask head head_loss_gradients['mask'] = _add_roi_mask_head( model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv ) if cfg.MODEL.KEYPOINTS_ON: # Add the keypoint head head_loss_gradients['keypoint'] = _add_roi_keypoint_head( model, add_roi_keypoint_head_func, blob_conv, dim_conv, spatial_scale_conv ) if model.train: loss_gradients = {} for lg in head_loss_gradients.values(): if lg is not None: loss_gradients.update(lg) return loss_gradients else: return None optim.build_data_parallel_model(model, _single_gpu_build_func) return model