diff --git a/tools/convert_pkl_to_pb.py b/tools/convert_pkl_to_pb.py index 27d4378ac..a553444bc 100644 --- a/tools/convert_pkl_to_pb.py +++ b/tools/convert_pkl_to_pb.py @@ -50,6 +50,7 @@ from detectron.utils.logging import setup_logging from detectron.utils.model_convert_utils import convert_op_in_proto from detectron.utils.model_convert_utils import op_filter +import detectron.utils.blob as blob_utils import detectron.core.test_engine as test_engine import detectron.utils.c2 as c2_utils import detectron.utils.model_convert_utils as mutils @@ -124,10 +125,41 @@ def unscope_name(name): def reset_names(names): - for i in range(0, len(names)): + for i in range(len(names)): names[i] = unscope_name(names[i]) +def convert_collect_and_distribute( + op, blobs, + roi_canonical_scale, + roi_canonical_level, + roi_max_level, + roi_min_level, + rpn_max_level, + rpn_min_level, + rpn_post_nms_topN, +): + print('Converting CollectAndDistributeFpnRpnProposals' + ' Python -> C++:\n{}'.format(op)) + assert op.name.startswith('CollectAndDistributeFpnRpnProposalsOp'), \ + 'Not valid CollectAndDistributeFpnRpnProposalsOp' + + inputs = [x for x in op.input] + ret = core.CreateOperator( + 'CollectAndDistributeFpnRpnProposals', + inputs, + list(op.output), + roi_canonical_scale=roi_canonical_scale, + roi_canonical_level=roi_canonical_level, + roi_max_level=roi_max_level, + roi_min_level=roi_min_level, + rpn_max_level=rpn_max_level, + rpn_min_level=rpn_min_level, + rpn_post_nms_topN=rpn_post_nms_topN, + ) + return ret + + def convert_gen_proposals( op, blobs, rpn_pre_nms_topN, @@ -136,19 +168,22 @@ def convert_gen_proposals( rpn_min_size, ): print('Converting GenerateProposals Python -> C++:\n{}'.format(op)) - assert op.name.startswith("GenerateProposalsOp"), "Not valid GenerateProposalsOp" + assert op.name.startswith('GenerateProposalsOp'), 'Not valid GenerateProposalsOp' - spatial_scale = mutils.get_op_arg_valf(op, "spatial_scale", None) + spatial_scale = mutils.get_op_arg_valf(op, 'spatial_scale', None) assert spatial_scale is not None + lvl = int(op.input[0][-1]) if op.input[0][-1].isdigit() else None + inputs = [x for x in op.input] - anchor_name = "anchor" + anchor_name = 'anchor{}'.format(lvl) if lvl else 'anchor' inputs.append(anchor_name) - blobs[anchor_name] = get_anchors(spatial_scale) + anchor_sizes = (cfg.FPN.RPN_ANCHOR_START_SIZE * 2.**(lvl - cfg.FPN.RPN_MIN_LEVEL),) if lvl else cfg.RPN.SIZES + blobs[anchor_name] = get_anchors(spatial_scale, anchor_sizes) print('anchors {}'.format(blobs[anchor_name])) ret = core.CreateOperator( - "GenerateProposals", + 'GenerateProposals', inputs, list(op.output), spatial_scale=spatial_scale, @@ -158,14 +193,13 @@ def convert_gen_proposals( min_size=rpn_min_size, correct_transform_coords=True, ) - return ret, anchor_name -def get_anchors(spatial_scale): +def get_anchors(spatial_scale, anchor_sizes): anchors = generate_anchors.generate_anchors( stride=1. / spatial_scale, - sizes=cfg.RPN.SIZES, + sizes=anchor_sizes, aspect_ratios=cfg.RPN.ASPECT_RATIOS).astype(np.float32) return anchors @@ -188,25 +222,64 @@ def convert_op_name(op): reset_names(op.output) return [op] - @op_filter(type="Python", inputs=['rpn_cls_probs', 'rpn_bbox_pred', 'im_info']) - def convert_gen_proposal(op_in): - gen_proposals_op, ext_input = convert_gen_proposals( - op_in, blobs, - rpn_min_size=float(cfg.TEST.RPN_MIN_SIZE), - rpn_post_nms_topN=cfg.TEST.RPN_POST_NMS_TOP_N, - rpn_pre_nms_topN=cfg.TEST.RPN_PRE_NMS_TOP_N, - rpn_nms_thresh=cfg.TEST.RPN_NMS_THRESH, - ) - net.external_input.extend([ext_input]) - return [gen_proposals_op] + @op_filter(type='Python') + def convert_python(op): + if op.name.startswith('GenerateProposalsOp'): + gen_proposals_op, ext_input = convert_gen_proposals( + op, blobs, + rpn_min_size=float(cfg.TEST.RPN_MIN_SIZE), + rpn_post_nms_topN=cfg.TEST.RPN_POST_NMS_TOP_N, + rpn_pre_nms_topN=cfg.TEST.RPN_PRE_NMS_TOP_N, + rpn_nms_thresh=cfg.TEST.RPN_NMS_THRESH, + ) + net.external_input.extend([ext_input]) + return [gen_proposals_op] + elif op.name.startswith('CollectAndDistributeFpnRpnProposalsOp'): + collect_dist_op = convert_collect_and_distribute( + op, blobs, + roi_canonical_scale=cfg.FPN.ROI_CANONICAL_SCALE, + roi_canonical_level=cfg.FPN.ROI_CANONICAL_LEVEL, + roi_max_level=cfg.FPN.ROI_MAX_LEVEL, + roi_min_level=cfg.FPN.ROI_MIN_LEVEL, + rpn_max_level=cfg.FPN.RPN_MAX_LEVEL, + rpn_min_level=cfg.FPN.RPN_MIN_LEVEL, + rpn_post_nms_topN=cfg.TEST.RPN_POST_NMS_TOP_N, + ) + return [collect_dist_op] + else: + raise ValueError('Failed to convert Python op {}'.format( + op.name)) + + # Only convert UpsampleNearest to ResizeNearest when converting to pb so that the existing models is unchanged + # https://github.com/facebookresearch/Detectron/pull/372#issuecomment-410248561 + @op_filter(type='UpsampleNearest') + def convert_upsample_nearest(op): + for arg in op.arg: + if arg.name == 'scale': + scale = arg.i + break + else: + raise KeyError('No attribute "scale" in UpsampleNearest op') + resize_nearest_op = core.CreateOperator('ResizeNearest', + list(op.input), + list(op.output), + name=op.name, + width_scale=float(scale), + height_scale=float(scale)) + return resize_nearest_op - @op_filter(input_has='rois') + @op_filter() def convert_rpn_rois(op): - for j in range(0, len(op.input)): + for j in range(len(op.input)): if op.input[j] == 'rois': print('Converting op {} input name: rois -> rpn_rois:\n{}'.format( op.type, op)) op.input[j] = 'rpn_rois' + for j in range(len(op.output)): + if op.output[j] == 'rois': + print('Converting op {} output name: rois -> rpn_rois:\n{}'.format( + op.type, op)) + op.output[j] = 'rpn_rois' return [op] @op_filter(type_in=['StopGradient', 'Alias']) @@ -214,10 +287,13 @@ def convert_remove_op(op): print('Removing op {}:\n{}'.format(op.type, op)) return [] + # We want to apply to all operators, including converted + # so run separately + convert_op_in_proto(net, convert_remove_op) + convert_op_in_proto(net, convert_upsample_nearest) + convert_op_in_proto(net, convert_python) convert_op_in_proto(net, convert_op_name) - convert_op_in_proto(net, [ - convert_gen_proposal, convert_rpn_rois, convert_remove_op - ]) + convert_op_in_proto(net, convert_rpn_rois) reset_names(net.external_input) reset_names(net.external_output) @@ -272,6 +348,7 @@ def convert_model_gpu(args, net, init_net): cdo_cpu = mutils.get_device_option_cpu() CPU_OPS = [ + ["CollectAndDistributeFpnRpnProposals", None], ["GenerateProposals", None], ["BBoxTransform", None], ["BoxWithNMSLimit", None], @@ -424,10 +501,8 @@ def _prepare_blobs( im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR) - blob = np.zeros([1, im.shape[0], im.shape[1], 3], dtype=np.float32) - blob[0, :, :, :] = im - channel_swap = (0, 3, 1, 2) # swap channel to (k, c, h, w) - blob = blob.transpose(channel_swap) + # Reuse code in blob_utils and fit FPN + blob = blob_utils.im_list_to_blob([im]) blobs = {} blobs['data'] = blob @@ -462,7 +537,7 @@ def run_model_pb(args, net, init_net, im, check_blobs): ) try: - workspace.RunNet(net.Proto().name) + workspace.RunNet(net) scores = workspace.FetchBlob('score_nms') classids = workspace.FetchBlob('class_nms') boxes = workspace.FetchBlob('bbox_nms') @@ -520,13 +595,16 @@ def main(): merge_cfg_from_list(args.opts) cfg.NUM_GPUS = 1 assert_and_infer_cfg() - logger.info('Conerting model with config:') + logger.info('Converting model with config:') logger.info(pprint.pformat(cfg)) - assert not cfg.MODEL.KEYPOINTS_ON, "Keypoint model not supported." - assert not cfg.MODEL.MASK_ON, "Mask model not supported." - assert not cfg.FPN.FPN_ON, "FPN not supported." - assert not cfg.RETINANET.RETINANET_ON, "RetinaNet model not supported." + # script will stop when it can't find an operator rather + # than stopping based on these flags + # + # assert not cfg.MODEL.KEYPOINTS_ON, "Keypoint model not supported." + # assert not cfg.MODEL.MASK_ON, "Mask model not supported." + # assert not cfg.FPN.FPN_ON, "FPN not supported." + # assert not cfg.RETINANET.RETINANET_ON, "RetinaNet model not supported." # load model from cfg model, blobs = load_model(args)