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
