def train_mask_rcnn(queue=None, imdb_name=None, init_model=None, solver=None, max_iters=None, cfg=None, rpn_file=None): """Train a Mask 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 cfg.TRAIN.IMS_PER_BATCH = 1 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 Mask R-CNN model_paths = train_net_multi_gpu(solver, roidb, output_dir, pretrained_model=init_model, max_iters=max_iters, gpus=cfg.GPU_ID) # Cleanup all but the final model # for i in model_paths[:-1]: # os.remove(i) mask_rcnn_model_path = model_paths[-1] # Send Mask R-CNN model path over the multiprocessing queue queue.put({'model_path': mask_rcnn_model_path})
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 Mask 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) 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_gpu(solver, roidb, output_dir, pretrained_model=init_model, max_iters=max_iters, gpus=cfg.GPU_ID) # 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})
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) gpu_id = args.gpu_id gpu_list = gpu_id.split(',') gpus = [int(i) for i in gpu_list] print('Using config:') pprint.pprint(cfg) if not args.randomize: # fix the random seeds (numpy and caffe) for reproducibility np.random.seed(cfg.RNG_SEED) # set up caffe imdb, roidb = combined_roidb(args.imdb_name) print '{:d} roidb entries'.format(len(roidb)) output_dir = get_output_dir(imdb) print 'Output will be saved to `{:s}`'.format(output_dir) train_net_multi_gpu(args.solver, roidb, output_dir, pretrained_model=args.pretrained_model, max_iter=args.max_iters, gpus=gpus)
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) gpu_id = args.gpu_id gpu_list = gpu_id.split(',') gpus = [int(i) for i in gpu_list] print('Using config:') pprint.pprint(cfg) if not args.randomize: # 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 imdb, roidb = combined_roidb(args.imdb_name) print '{:d} roidb entries'.format(len(roidb)) output_dir = get_output_dir(imdb) print 'Output will be saved to `{:s}`'.format(output_dir) train_net_multi_gpu(args.solver, roidb, output_dir, pretrained_model=args.pretrained_model, max_iter=args.max_iters, gpus=gpus)