train_faster_rcnn_alt_opt.py 文件源码

def rpn_generate(gpus, queue=None, imdb_name=None, rpn_model_path=None, cfg=None,
                 rpn_test_prototxt=None):
    """Use a trained RPN to generate proposals.
    """
    def rpn_generate_signle_gpu(rank):
        cfg.GPU_ID=gpus[rank]

        print('Using config:')
        pprint.pprint(cfg)

        import caffe
        np.random.seed(cfg.RNG_SEED)
        caffe.set_random_seed(cfg.RNG_SEED)
        # set up caffe
        caffe.set_mode_gpu()
        caffe.set_device(cfg.GPU_ID)

        # Load RPN and configure output directory
        rpn_net = caffe.Net(rpn_test_prototxt, rpn_model_path, caffe.TEST)

        # Generate proposals on the imdb
        rpn_proposals = imdb_proposals(rpn_net, imdb, rank, len(gpus), output_dir)


    cfg.TEST.RPN_PRE_NMS_TOP_N = -1     # no pre NMS filtering
    cfg.TEST.RPN_POST_NMS_TOP_N = 2000  # limit top boxes after NMS

    print 'RPN model: {}'.format(rpn_model_path)
    imdb = get_imdb(imdb_name)

    output_dir = os.path.join(get_output_dir(imdb), "proposals")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    print 'Output will be saved to `{:s}`'.format(output_dir)
    # NOTE: the matlab implementation computes proposals on flipped images, too.
        # We compute them on the image once and then flip the already computed
        # proposals. This might cause a minor loss in mAP (less proposal jittering).
    print 'Loaded dataset `{:s}` for proposal generation'.format(imdb.name)

    procs=[]
    for rank in range(len(gpus)):
        p = mp.Process(target=rpn_generate_signle_gpu,
                    args=(rank, ))
        p.daemon = True
        p.start()
        procs.append(p)
    for p in procs:
        p.join()
    queue.put({'proposal_path': output_dir})