def convert_net()

in tools/convert_pkl_to_pb.py [0:0]

• 86 lines of code
• 11 McCabe index (conditional complexity)

def convert_net(args, net, blobs):
    @op_filter()
    def convert_op_name(op):
        if args.device != "gpu":
            if op.engine != "DEPTHWISE_3x3":
                op.engine = ""
            op.device_option.CopyFrom(caffe2_pb2.DeviceOption())
        reset_names(op.input)
        reset_names(op.output)
        return [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()
    def convert_rpn_rois(op):
        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"])
    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_rpn_rois)

    reset_names(net.external_input)
    reset_names(net.external_output)

    reset_blob_names(blobs)