py-faster-rcnn详解（1）——train_faster_rcnn_alt_opt.py接口说明

  py-faster-rcnn使用experiments/scripts/faster_rcnn_alt_opt.sh脚本启动训练过程。在faster_rcnn_alt_opt.sh脚本中，调用tools/train_faster_rcnn_alt_opt.py接口

1、parse_args()
    #传递参数 
2、 get_roidb(imdb_name, rpn_file= None)
  #获取roidb与lmdb 
3、get_solvers(net_name)  #
获取solvers 
4、_init_caffe(cfg)
  #根据config加载caffe对象 
5、train_rpn(queue, imdb_name, init_model, solver=None,
 max_iters, cfg)
  #训练RPN网络 
6、 rpn_generate(queue, imdb_name, rpn_model_path, cfg,
 rpn_test_prototxt)
   #用RPN网络生产proposal 
7、 train_fast_rcnn(queue, imdb_name, init_model, solver,
 max_iters, cfg, rpn_file)
  #用RPN网络产生的proposal来训练fast_rcnn网络 
8、mian()  #详细的描述了四个训练步骤

get_roidb

  imdb根据imdb_name(默认是“voc_2007_trainval)来获取；

  roidb是通过get_training_roidb来获取的，roidb是imdb的成员变量，包含了训练集图片中框出的每个区域。

def get_roidb(imdb_name, rpn_file=None):
 imdb = get_imdb(imdb_name)
 print 'Loaded dataset {:s} for training'.format(imdb.name)
 imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD)
 print 'Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD)
 if rpn_file is not None:
     imdb.config['rpn_file'] = rpn_file
 roidb = get_training_roidb(imdb)
 return roidb, imdb
   
   
1
2
3
4
5
6
7
8
9
   
   
1
2
3
4
5
6
7
8
9

get_solvers

  在models/pascal_voc/netname/faster_rcnn_alt_opt文件夹下有stage1_rpn_solver60k80k.pt等不同阶段所对应的solver文件，并定义了各个阶段最大迭代次数，这里支持的net_name有 VGG16、VGG_CNN_M_1024、ZF三种

def get\_solvers(net\_name):
 # Faster R-CNN Alternating Optimization
 n = 'faster_rcnn_alt_opt'
 # Solver for each training stage
 solvers = [[net_name, n, 'stage1_rpn_solver60k80k.pt'],
            [net_name, n, 'stage1_fast_rcnn_solver30k40k.pt'],
            [net_name, n, 'stage2_rpn_solver60k80k.pt'],
            [net_name, n, 'stage2_fast_rcnn_solver30k40k.pt']]
 solvers = [os.path.join(cfg.MODELS_DIR, *s) for s in solvers]
 # Iterations for each training stage
 max_iters = [80000, 40000, 80000, 40000]
 # max_iters = [100, 100, 100, 100]
 # Test prototxt for the RPN
 rpn_test_prototxt = os.path.join(
     cfg.MODELS_DIR, net_name, n, 'rpn_test.pt')
 return solvers, max_iters, rpn_test_prototxt
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

_init_caffe

  初始化caffe对象，一是初始化随机种子，二是设置GPU。

def _init_caffe(cfg):
    """Initialize pycaffe in a training process.
    """
    import caffe
    # fix the random seeds (numpy and caffe) for reproducibility
    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)

   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

train_rpn

 训练rpn网络。首先这里初始化了一些配置；随后做了一些训练的准备工作，初始化模型，确定路径等；然后开始正式训练。

 def train_rpn(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 box regression
 # 用ground_truth作为proposal
 cfg.TRAIN.PROPOSAL_METHOD = 'gt'
 # 每个minibatch训练用到的Image数
 cfg.TRAIN.IMS_PER_BATCH = 1
 print 'Init model: {}'.format(init_model)
 print('Using config:')
 pprint.pprint(cfg)
 #初始化caffe
 import caffe
 _init_caffe(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(solver, roidb, output_dir,
                         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})
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

rpn_generate

  根据输入的数据与模型与prototxt产生proposal，可作为下一步的训练所用。 

def rpn_generate(queue=None, imdb\_name=None, rpn\_model\_path=None, cfg=None,rpn_test_prototxt=None):
"""Use a trained RPN to generate proposals.
"""
    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)
    print('Using config:')
    pprint.pprint(cfg)

    import caffe
    _init_caffe(cfg)

    # 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).
    imdb = get_imdb(imdb_name)
    print 'Loaded dataset {:s} for proposal generation'.format(imdb.name)

    # Load RPN and configure output directory
    rpn_net = caffe.Net(rpn_test_prototxt, rpn_model_path, caffe.TEST)
    output_dir = get_output_dir(imdb)
    print 'Output will be saved to {:s}'.format(output_dir)
    # Generate proposals on the imdb
    rpn_proposals = imdb_proposals(rpn_net, imdb)
    # Write proposals to disk and send the proposal file path through the
    # multiprocessing queue
    rpn_net_name = os.path.splitext(os.path.basename(rpn_model_path))[0]
    rpn_proposals_path = os.path.join(
    output_dir, rpn_net_name + '_proposals.pkl')
    with open(rpn_proposals_path, 'wb') as f:
    cPickle.dump(rpn_proposals, f, cPickle.HIGHEST_PROTOCOL)
    print 'Wrote RPN proposals to {}'.format(rpn_proposals_path)
    queue.put({'proposal_path': rpn_proposals_path})

   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

train_fast_rcnn

  训练fast-rcnn，首先它将产生roidb的方法设置成rpn_roidb；接下来就是准备一些参数、路径等等，用于送入网络训练，最后保存模型。

def train_fast_rcnn(queue=None, imdb_name=None, init_model=None, solver=None,
max_iters=None, cfg=None, rpn_file=None):
"""Train a Fast R-CNN using proposals generated by an RPN.
"""
    cfg.TRAIN.HAS_RPN = False           # not generating prosals on-the-fly
    cfg.TRAIN.PROPOSAL_METHOD = 'rpn'   # use pre-computed RPN proposals instead
    #每个mini-batch包含两张图片，以及他们proposal的roi区域
    cfg.TRAIN.IMS_PER_BATCH = 2
    print 'Init model: {}'.format(init_model)
    print 'RPN proposals: {}'.format(rpn_file)
    print('Using config:')
    pprint.pprint(cfg)

    import caffe
    _init_caffe(cfg)

    roidb, imdb = get_roidb(imdb_name, rpn_file=rpn_file)
    output_dir = get_output_dir(imdb)
    print 'Output will be saved to {:s}'.format(output_dir)
    # Train Fast R-CNN
    model_paths = train_net(solver, roidb, output_dir,
    pretrained_model=init_model,
    max_iters=max_iters)
    # Cleanup all but the final model
    for i in model_paths[:-1]:
    os.remove(i)
    fast_rcnn_model_path = model_paths[-1]
    # Send Fast R-CNN model path over the multiprocessing queue
    queue.put({'model_path': fast_rcnn_model_path})




   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

if __name__ == '__main__':
args = parse_args()

print('Called with args:')
print(args)
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
cfg.GPU_ID = args.gpu_id

# --------------------------------------------------------------------------
# Pycaffe doesn't reliably free GPU memory when instantiated nets are
# discarded (e.g. "del net" in Python code). To work around this issue, each
# training stage is executed in a separate process using
# multiprocessing.Process.
# --------------------------------------------------------------------------
# 这里不同的训练stage将通过多线程来提高效率，mp_queue是进程间用于通讯的数据结构
# queue for communicated results between processes
mp_queue = mp.Queue()
# 获取 solves, iters, etc. for each training stage
solvers, max_iters, rpn_test_prototxt = get_solvers(args.net_name)

 # 这一步是用imageNet的模型初始化，然后训练rpn网络（整个训练过程可以参考作者的论文） 
print
print 'Stage 1 RPN, init from ImageNet model'
print
# 第一个stage开始啦~~
cfg.TRAIN.SNAPSHOT_INFIX = 'stage1'
mp_kwargs = dict(
queue=mp_queue,
imdb_name=args.imdb_name,
init_model=args.pretrained_model,
solver=solvers[0],
max_iters=max_iters[0],
cfg=cfg)
p = mp.Process(target=train_rpn, kwargs=mp_kwargs)
p.start()
rpn_stage1_out = mp_queue.get()
p.join()

# 这一步是利用上一步训练好的rpn网络，产生proposals供后面使用 
print
print 'Stage 1 RPN, generate proposals'
print
mp_kwargs = dict(
queue=mp_queue,
imdb_name=args.imdb_name,
rpn_model_path=str(rpn_stage1_out['model_path']),
cfg=cfg,
rpn_test_prototxt=rpn_test_prototxt)
p = mp.Process(target=rpn_generate, kwargs=mp_kwargs)
p.start()
rpn_stage1_out['proposal_path'] = mp_queue.get()['proposal_path']
p.join()

#这一步是再次用imageNet的模型初始化前5层卷积层，然后用上一步得到的proposals训练检测网络 
print
print 'Stage 1 Fast R-CNN using RPN proposals, init from ImageNet model'
print
cfg.TRAIN.SNAPSHOT_INFIX = 'stage1'
mp_kwargs = dict(
queue=mp_queue,
imdb_name=args.imdb_name,
init_model=args.pretrained_model,
solver=solvers[1],
max_iters=max_iters[1],
cfg=cfg,
rpn_file=rpn_stage1_out['proposal_path'])
p = mp.Process(target=train_fast_rcnn, kwargs=mp_kwargs)
p.start()
fast_rcnn_stage1_out = mp_queue.get()
p.join()

#这一步固定上一步训练好的前五层卷积层，再次训练RPN，这样就得到最终RPN网络的参数了 
print
print 'Stage 2 RPN, init from stage 1 Fast R-CNN model'
print
cfg.TRAIN.SNAPSHOT_INFIX = 'stage2'
mp_kwargs = dict(
queue=mp_queue,
imdb_name=args.imdb_name,
init_model=str(fast_rcnn_stage1_out['model_path']),
solver=solvers[2],
max_iters=max_iters[2],
cfg=cfg)
p = mp.Process(target=train_rpn, kwargs=mp_kwargs)
p.start()
rpn_stage2_out = mp_queue.get()
p.join()

 #利用最终确定的RPN网络产生proposals  
print
print 'Stage 2 RPN, generate proposals'
print
mp_kwargs = dict(
queue=mp_queue,
imdb_name=args.imdb_name,
rpn_model_path=str(rpn_stage2_out['model_path']),
cfg=cfg,
rpn_test_prototxt=rpn_test_prototxt)
p = mp.Process(target=rpn_generate, kwargs=mp_kwargs)
p.start()
rpn_stage2_out['proposal_path'] = mp_queue.get()['proposal_path']
p.join()

#利用上一步产生的proposals，训练出最终的检测网络  
print
print 'Stage 2 Fast R-CNN, init from stage 2 RPN R-CNN model'
print
cfg.TRAIN.SNAPSHOT_INFIX = 'stage2'
mp_kwargs = dict(
queue=mp_queue,
imdb_name=args.imdb_name,
init_model=str(rpn_stage2_out['model_path']),
solver=solvers[3],
max_iters=max_iters[3],
cfg=cfg,
rpn_file=rpn_stage2_out['proposal_path'])
p = mp.Process(target=train_fast_rcnn, kwargs=mp_kwargs)
p.start()
fast_rcnn_stage2_out = mp_queue.get()
p.join()
#输出最终模型
# Create final model (just a copy of the last stage)
final_path = os.path.join(
os.path.dirname(fast_rcnn_stage2_out['model_path']),
args.net_name + '_faster_rcnn_final.caffemodel')
print 'cp {} -> {}'.format(
fast_rcnn_stage2_out['model_path'], final_path)
shutil.copy(fast_rcnn_stage2_out['model_path'], final_path)
print 'Final model: {}'.format(final_path)
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
   
   
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134