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_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)
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
optim.build_data_parallel_model(model, _single_gpu_build_func)
return model
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)
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
optim.build_data_parallel_model(model, _single_gpu_build_func)
return model
def build_generic_retinanet_model_dissstillation(
model, add_conv_body_func, 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."""
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
optim.build_data_parallel_model(model, _single_gpu_build_func)
return model
def build_generic_rc_model(model,
add_conv_body_func,
add_roi_box_head_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."""
blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model)
if not model.train:
model.conv_body_net = model.net.Clone('conv_body_net')
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 = {}
# Add the Fast R-CNN head
head_loss_gradients['box'] = _add_fast_rcnn_head_class_only(
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
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
)
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_static_memory_model(model,
add_conv_body_func,
add_roi_box_head_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."""
blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model)
if not model.train:
model.conv_body_net = model.net.Clone('conv_body_net')
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)
# break the fast rcnn head down
blob_frcn, dim_frcn = add_roi_box_head_func(model, blob_conv, dim_conv,
spatial_scale_conv)
fast_rcnn_heads.add_fast_rcnn_outputs_class_only(
model, blob_frcn, dim_frcn)
head_loss_gradients = {}
if model.train:
head_loss_gradients[
'base'] = fast_rcnn_heads.add_fast_rcnn_losses_class_only(
model)
image_blob_name = core.ScopedName('data')
rois_name = core.ScopedName('rois')
if 'gpu_0' in rois_name:
model.AddSummaryImageBoxes(image_blob_name, rois_name)
if cfg.FPN.FPN_ON:
blob_conv = [
model.StopGradient(bc,
c2_utils.UnscopeGPUName(bc._name + '_nb'))
for bc in blob_conv
]
else:
blob_conv = model.StopGradient(
blob_conv, c2_utils.UnscopeGPUName(blob_conv._name + '_nb'))
cls_score = u'cls_score'
cls_score_base = model.StopGradient(cls_score, cls_score + '_nb')
cls_prob = u'cls_prob'
cls_prob_base = core.ScopedBlobReference(cls_prob)
# cls_prob_base = model.StopGradient(cls_prob, cls_prob + '_nb')
mem = region_memory_model.init(model)
cls_score_list = [cls_score_base]
norm = region_memory_model.init_normalizer(model)
if 'gpu_0' in mem._name:
model.AddSummaryMem(mem._name)
if cfg.MEM.AT_MIN:
cls_attend_list = []
else:
cls_attend_list = [
region_memory_model.init_attenton_prediction(model, mem)
]
cls_score = cls_score_base
cls_prob = cls_prob_base
reuse = False
conv_crop = region_memory_model._roi_align(model, blob_conv,
spatial_scale_conv)
conv_crop_nb = model.StopGradient(
conv_crop, c2_utils.UnscopeGPUName(conv_crop._name + '_nb'))
norm_crop = region_memory_model._norm_roi_align(model, norm)
norm_diff = model.InvRoIAlign(core.ScopedBlobReference('rois'), norm,
norm_crop)
if 'gpu_0' in norm_diff._name:
model.AddSummaryMem(norm_diff._name)
for iter in range(1, cfg.MEM.ITER + 1):
mem = region_memory_model.update(model,
mem,
norm_diff,
conv_crop_nb,
dim_conv,
cls_score,
cls_prob,
iter,
reuse=reuse)
if 'gpu_0' in mem._name:
model.AddSummaryMem(mem._name)
# for testing, return cls_prob
cls_score, cls_prob, cls_attend = region_memory_model.prediction(
model, mem, cls_score_base, iter, reuse=reuse)
# for training, it will get cls_prob when getting the loss
if model.train:
name = 'mem_%02d' % iter
head_loss_gradients[
name], cls_prob = region_memory_model.add_loss(
model, cls_score, cfg.MEM.WEIGHT)
cls_score = model.StopGradient(
cls_score, c2_utils.UnscopeGPUName(cls_score._name + '_nb'))
# cls_prob = model.StopGradient(cls_prob, c2_utils.UnscopeGPUName(cls_prob._name + '_nb'))
cls_score_list.append(cls_score)
cls_attend_list.append(cls_attend)
reuse = True
cls_score_final = region_memory_model.combine(model, cls_score_list,
cls_attend_list)
if model.train:
head_loss_gradients[
'final'], cls_prob_final = region_memory_model.add_loss(
model, cls_score_final, cfg.MEM.WEIGHT_FINAL)
loss_gradients = {}
for lg in head_loss_gradients.values():
if lg is not None:
loss_gradients.update(lg)
return loss_gradients
else:
cls_prob_final = region_memory_model.add_final_prob(
model, cls_score_final)
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
)
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
