r-cnn學習（四）：train_faster_rcnn_alt_opt.py源碼學習

論文看的雲里霧里，希望通過閱讀其代碼來進一步了解。

參考：http://blog.csdn.net/sloanqin/article/details/51525692

 首先是./tools/train_faster_rcnn_alt_opt.py，通過其main函數了解整個訓練流程。

if __name__ == '__main__': #建議讀者調試這個函數，進去看看每個變量是怎么回事 
    args = parse_args() #解析系統傳入的argv參數，解析完放到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 # cfg是一個詞典（edict）數據結構，從faster-rcnn.config引入的 
  
    # -------------------------------------------------------------------------- 
    # 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. #這里說的要使用多進程,因為在pycaffe中當某個網絡被discard后，不能可靠保證釋放內存資源；進程關閉后資源自然會釋放 
    # -------------------------------------------------------------------------- 
  
    # queue for communicated results between processes 
    mp_queue = mp.Queue() #mp指的是multiprocessing庫，所以這里返回了一個用於多進程通信的隊列對象 
    # solves, iters, etc. for each training stage 
    solvers, max_iters, rpn_test_prototxt = get_solvers(args.net_name) #這里返回了solvers的路徑，maxiters的值，rpn_test_prototxt的路徑 
  
    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    print 'Stage 1 RPN, init from ImageNet model'  
    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    # 這一步是用imageNet的模型初始化，然后訓練rpn網絡（整個訓練過程可以參考作者的論文） 
    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) # 這里把該階段需要的參數都放到這里來了，即函數train_rpn的輸入參數 
    p = mp.Process(target=train_rpn, kwargs=mp_kwargs) # 顯然，這里准備啟動一個新進程，調用函數train_rpn，傳入參數kwargs，所以我們進入train_rpn函數看看是如何工作的 
 p.start() rpn_stage1_out = mp_queue.get() p.join() print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    print 'Stage 1 RPN, generate proposals'  
    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    # 這一步是利用上一步訓練好的rpn網絡，產生proposals供后面使用 
    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() print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    print 'Stage 1 Fast R-CNN using RPN proposals, init from ImageNet model'  
    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    #這一步是再次用imageNet的模型初始化前5層卷積層，然后用上一步得到的proposals訓練檢測網絡 
    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() print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    print 'Stage 2 RPN, init from stage 1 Fast R-CNN model'  
    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    #這一步固定上一步訓練好的前五層卷積層，再次訓練RPN，這樣就得到最終RPN網絡的參數了 
    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() print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    print 'Stage 2 RPN, generate proposals'  
    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    #利用最終確定的RPN網絡產生proposals 
    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() print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    print 'Stage 2 Fast R-CNN, init from stage 2 RPN R-CNN model'  
    print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
    #利用上一步產生的proposals，訓練出最終的檢測網絡 
    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)  

通過上面的代碼可以看出，整個迭代過程分為四步（參考論文）。其中后面兩步固定共享卷積

層，只對RPN和fc層進行微調。

 

接着看看每一步是怎樣的。

首先是train_rpn。從代碼看出，這個函數的主要任務是，配置參數，准備數據集，

傳入第一階段的solver，調用train_net訓練模型並將結果返回。

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. """  
    #首先進來后繼續配置了一些cfg這個對象的一些參數 
    # 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  
    print 'Init model: {}'.format(init_model) #格式化輸出字符串 
    print('Using config:') pprint.pprint(cfg) import caffe _init_caffe(cfg) #這里是關鍵，准備數據集，我們在debug的時候可以發現，imdb是一個類，而roidb是該類的一個成員 
    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) #這個solver傳入的是./models/pascal_voc/ZF/faster_rcnn_alt_opt/stage1_rpn_solver60k80k.pt 
    model_paths = train_net(solver, roidb, output_dir, pretrained_model=init_model, max_iters=max_iters) #進入train_net函數，看訓練如何實現的 
    # 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}) #通過隊列將該進程運行的模型結果的路徑返回  

 

順着train_rpn，查看train_net函數，該函數位於：./lib/fast_rcnn/train.py文件中

調用該文件中定義的類SolverWrapper的構造函數，返回該類的一個對象sw，然后調用了sw的train_model方法進行訓練，

傳入參數，搭建caffe的網絡結構，用預訓練模型完成初始化，整個過程在構造函數中完成。

 

"""Train a Fast R-CNN network."""  
  
import caffe from fast_rcnn.config import cfg import roi_data_layer.roidb as rdl_roidb from utils.timer import Timer import numpy as np import os from caffe.proto import caffe_pb2 import google.protobuf as pb2 class SolverWrapper(object): """A simple wrapper around Caffe's solver. This wrapper gives us control over he snapshotting process, which we use to unnormalize the learned bounding-box regression weights. """  
  
    #這就是SolverWrapper的構造函數 
    def __init__(self, solver_prototxt, roidb, output_dir, pretrained_model=None): """Initialize the SolverWrapper.""" self.output_dir = output_dir if (cfg.TRAIN.HAS_RPN and cfg.TRAIN.BBOX_REG and cfg.TRAIN.BBOX_NORMALIZE_TARGETS): # RPN can only use precomputed normalization because there are no 
            # fixed statistics to compute a priori 
            assert cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED if cfg.TRAIN.BBOX_REG: print 'Computing bounding-box regression targets...' self.bbox_means, self.bbox_stds = \ rdl_roidb.add_bbox_regression_targets(roidb) print 'done'  
  
        # 這句話調用了caffe的SGDSolver，這個是caffe在C++中實現的一個類，用來進行隨機梯度下降優化，該類根據solver_prototxt中定義的網絡和求解參數，完成網絡 
        # 初始化，然后返回類SGDSolver的一個實例，關於該類的設計可以參考caffe的網站：http://caffe.berkeleyvision.org/doxygen/classcaffe_1_1SGDSolver.html 
        # 然后作者把該對象作為SolverWrapper的一個成員，命名為solver 
        self.solver = caffe.SGDSolver(solver_prototxt) if pretrained_model is not None: print ('Loading pretrained model '  
                   'weights from {:s}').format(pretrained_model) self.solver.net.copy_from(pretrained_model)#這句話完成對網絡的初始化 
 self.solver_param = caffe_pb2.SolverParameter() with open(solver_prototxt, 'rt') as f: pb2.text_format.Merge(f.read(), self.solver_param)#這句話應該是設置了self.solver_param這個成員的參數 
 self.solver.net.layers[0].set_roidb(roidb)#這句話傳入訓練的數據：roidb 
  
    def snapshot(self): """Take a snapshot of the network after unnormalizing the learned bounding-box regression weights. This enables easy use at test-time. """ net = self.solver.net scale_bbox_params = (cfg.TRAIN.BBOX_REG and cfg.TRAIN.BBOX_NORMALIZE_TARGETS and net.params.has_key('bbox_pred')) if scale_bbox_params: # save original values 
            orig_0 = net.params['bbox_pred'][0].data.copy() orig_1 = net.params['bbox_pred'][1].data.copy() # scale and shift with bbox reg unnormalization; then save snapshot 
            net.params['bbox_pred'][0].data[...] = \ (net.params['bbox_pred'][0].data * self.bbox_stds[:, np.newaxis]) net.params['bbox_pred'][1].data[...] = \ (net.params['bbox_pred'][1].data * self.bbox_stds + self.bbox_means) infix = ('_' + cfg.TRAIN.SNAPSHOT_INFIX if cfg.TRAIN.SNAPSHOT_INFIX != '' else '') filename = (self.solver_param.snapshot_prefix + infix +  
                    '_iter_{:d}'.format(self.solver.iter) + '.caffemodel') filename = os.path.join(self.output_dir, filename) net.save(str(filename)) print 'Wrote snapshot to: {:s}'.format(filename) if scale_bbox_params: # restore net to original state 
            net.params['bbox_pred'][0].data[...] = orig_0 net.params['bbox_pred'][1].data[...] = orig_1 return filename def train_model(self, max_iters): """Network training loop.""" last_snapshot_iter = -1 timer = Timer() model_paths = [] while self.solver.iter < max_iters: # Make one SGD update 
            timer.tic()#作者測量一次迭代花的時間 
            self.solver.step(1)# 做一次梯度下降優化 
 timer.toc() if self.solver.iter % (10 * self.solver_param.display) == 0: print 'speed: {:.3f}s / iter'.format(timer.average_time) if self.solver.iter % cfg.TRAIN.SNAPSHOT_ITERS == 0: last_snapshot_iter = self.solver.iter model_paths.append(self.snapshot()) if last_snapshot_iter != self.solver.iter: model_paths.append(self.snapshot()) return model_paths def get_training_roidb(imdb): """Returns a roidb (Region of Interest database) for use in training."""  
    if cfg.TRAIN.USE_FLIPPED: print 'Appending horizontally-flipped training examples...' imdb.append_flipped_images() print 'done'  
  
    print 'Preparing training data...' rdl_roidb.prepare_roidb(imdb) print 'done'  
  
    return imdb.roidb def filter_roidb(roidb): """Remove roidb entries that have no usable RoIs."""  
    #判斷是否是有效roidb
    def is_valid(entry): # Valid images have: 
        # (1) At least one foreground RoI OR 
        # (2) At least one background RoI 
        overlaps = entry['max_overlaps'] # find boxes with sufficient overlap 
        fg_inds = np.where(overlaps >= cfg.TRAIN.FG_THRESH)[0]#大於某個閾值為前景 # Select background RoIs as those within [BG_THRESH_LO, BG_THRESH_HI) 
        bg_inds = np.where((overlaps < cfg.TRAIN.BG_THRESH_HI) & #在某兩個閾值之間為背景 (overlaps >= cfg.TRAIN.BG_THRESH_LO))[0] # image is only valid if such boxes exist 
        valid = len(fg_inds) > 0 or len(bg_inds) > 0#要么為前景，要么為背景，則為有效roidb return valid num = len(roidb) filtered_roidb = [entry for entry in roidb if is_valid(entry)] num_after = len(filtered_roidb) print 'Filtered {} roidb entries: {} -> {}'.format(num - num_after, num, num_after) return filtered_roidb # 該函數先是調用了該文件中定義的類SolverWrapper的構造函數，返回了該類的一個對象sw，然后調用了sw的train_model方法進行訓練  # 傳入參數，搭建caffe的網絡結構，用預訓練模型完成初始化，這些過程就是在該構造函數中實現的，進入這個構造函數看看 
def train_net(solver_prototxt, roidb, output_dir, pretrained_model=None, max_iters=40000): """Train a Fast R-CNN network.""" roidb = filter_roidb(roidb)#刪除一些不滿足要求的輸入圖片 
    sw = SolverWrapper(solver_prototxt, roidb, output_dir, pretrained_model=pretrained_model)#調用構造函數 
  
    print 'Solving...' model_paths = sw.train_model(max_iters)#開始訓練模型 
    print 'done solving'  
    return model_paths