train_rfcn_alt_opt_5stage.py 文件源码

def get_solvers(imdb_name, net_name, model_name):
    # R-FCN Alternating Optimization
    # Solver for each training stage
    if imdb_name.startswith('coco'):
        solvers = [[net_name, model_name, 'stage1_rpn_solver360k480k.pt'],
                   [net_name, model_name, 'stage1_rfcn_ohem_solver360k480k.pt'],
                   [net_name, model_name, 'stage2_rpn_solver360k480k.pt'],
                   [net_name, model_name, 'stage2_rfcn_ohem_solver360k480k.pt'],
                   [net_name, model_name, 'stage3_rpn_solver360k480k.pt']]
        solvers = [os.path.join('.', 'models', 'coco', *s) for s in solvers]
        # Iterations for each training stage
        max_iters = [480000, 480000, 480000, 480000, 480000]
        # Test prototxt for the RPN
        rpn_test_prototxt = os.path.join(
            '.', 'models', 'coco', net_name, model_name, 'rpn_test.pt')
    else:
        solvers = [[net_name, model_name, 'stage1_rpn_solver60k80k.pt'],
                   [net_name, model_name, 'stage1_rfcn_ohem_solver80k120k.pt'],
                   [net_name, model_name, 'stage2_rpn_solver60k80k.pt'],
                   [net_name, model_name, 'stage2_rfcn_ohem_solver80k120k.pt'],
                   [net_name, model_name, 'stage3_rpn_solver60k80k.pt']]
        solvers = [os.path.join(cfg.MODELS_DIR, *s) for s in solvers]
        # Iterations for each training stage
        max_iters = [80000, 120000, 80000, 120000, 80000]
        # Test prototxt for the RPN
        rpn_test_prototxt = os.path.join(
            cfg.MODELS_DIR, net_name, model_name, 'rpn_test.pt')
    return solvers, max_iters, rpn_test_prototxt