train_faster_rcnn_alt_opt.py 文件源码

def train_rpn(gpus, queue=None, imdb_name=None, init_model=None, solver=None,
              max_iters=None, cfg=None):
    """Train a Region Proposal Network in a separate training process.
    """

    # Not using any proposals, just ground-truth boxes
    cfg.TRAIN.HAS_RPN = True
    cfg.TRAIN.BBOX_REG = False  # applies only to Fast R-CNN bbox regression
    cfg.TRAIN.PROPOSAL_METHOD = 'gt'
    cfg.TRAIN.IMS_PER_BATCH = 1
    cfg.TRAIN.REAL_BATCH_SIZE = 8
    cfg.TRAIN.VAL_PER_BATCH_SIZE = 1
    np.random.seed(cfg.RNG_SEED)
    print 'Init model: {}'.format(init_model)
    print('Using config:')
    pprint.pprint(cfg)
    roidb, imdb = get_roidb(imdb_name)
    print 'roidb len: {}'.format(len(roidb))
    output_dir = get_output_dir(imdb)
    print 'Output will be saved to `{:s}`'.format(output_dir)

    model_paths = train_net_multi_gpus(solver, roidb, output_dir, gpus,
                            pretrained_model=init_model,
                            max_iters=max_iters)
    # Cleanup all but the final model
    # for i in model_paths[:-1]:
    #     os.remove(i)
    rpn_model_path = model_paths[-1]
    # Send final model path through the multiprocessing queue
    queue.put({'model_path': rpn_model_path})