def infer(): place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) class_nums = cfg.class_num model = model_builder.RRPN(add_conv_body_func=resnet.ResNet(), add_roi_box_head_func=resnet.ResNetC5(), use_pyreader=False, mode='infer') startup_prog = fluid.Program() infer_prog = fluid.Program() with fluid.program_guard(infer_prog, startup_prog): with fluid.unique_name.guard(): model.build_model() pred_boxes = model.eval_bbox_out() infer_prog = infer_prog.clone(True) exe.run(startup_prog) fluid.load(infer_prog, cfg.pretrained_model, exe) infer_reader = reader.infer(cfg.image_path) data_loader = model.data_loader data_loader.set_sample_list_generator(infer_reader, places=place) fetch_list = [pred_boxes] imgs = os.listdir(cfg.image_path) imgs.sort() for i, data in enumerate(data_loader()): result = exe.run(infer_prog, fetch_list=[v.name for v in fetch_list], feed=data, return_numpy=False) nmsed_out = result[0] im_info = np.array(data[0]['im_info'])[0] im_scale = im_info[2] outs = np.array(nmsed_out) draw_bounding_box_on_image(cfg.image_path, imgs[i], outs, im_scale, cfg.draw_threshold)
def eval(): place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) class_nums = cfg.class_num model = model_builder.RRPN(add_conv_body_func=resnet.ResNet(), add_roi_box_head_func=resnet.ResNetC5(), use_pyreader=False, mode='val') startup_prog = fluid.Program() infer_prog = fluid.Program() with fluid.program_guard(infer_prog, startup_prog): with fluid.unique_name.guard(): model.build_model() pred_boxes = model.eval_bbox_out() infer_prog = infer_prog.clone(True) exe.run(startup_prog) fluid.load(infer_prog, cfg.pretrained_model, exe) test_reader = reader.test(1) data_loader = model.data_loader data_loader.set_sample_list_generator(test_reader, places=place) fetch_list = [pred_boxes] res_list = [] keys = [ 'bbox', 'gt_box', 'gt_class', 'is_crowed', 'im_info', 'im_id', 'is_difficult' ] for i, data in enumerate(data_loader()): result = exe.run(infer_prog, fetch_list=[v.name for v in fetch_list], feed=data, return_numpy=False) pred_boxes_v = result[0] nmsed_out = pred_boxes_v outs = np.array(nmsed_out) res = get_key_dict(outs, data[0], keys) res_list.append(res) if i % 50 == 0: logger.info('test_iter {}'.format(i)) icdar_eval(res_list)
def train(): learning_rate = cfg.learning_rate #image_shape = [-1, 3, cfg.TRAIN.max_size, cfg.TRAIN.max_size] devices_num = get_device_num() total_batch_size = devices_num * cfg.TRAIN.im_per_batch use_random = True startup_prog = fluid.Program() train_prog = fluid.Program() with fluid.program_guard(train_prog, startup_prog): with fluid.unique_name.guard(): model = model_builder.RRPN(add_conv_body_func=resnet.ResNet(), add_roi_box_head_func=resnet.ResNetC5(), use_pyreader=cfg.use_pyreader, use_random=use_random) model.build_model() losses, keys, rpn_rois = model.loss() loss = losses[0] fetch_list = losses boundaries = cfg.lr_steps gamma = cfg.lr_gamma step_num = len(cfg.lr_steps) values = [learning_rate * (gamma**i) for i in range(step_num + 1)] start_lr = learning_rate * cfg.start_factor lr = fluid.layers.piecewise_decay(boundaries, values) lr = fluid.layers.linear_lr_warmup(lr, cfg.warm_up_iter, start_lr, learning_rate) optimizer = fluid.optimizer.Momentum( learning_rate=lr, regularization=fluid.regularizer.L2Decay(cfg.weight_decay), momentum=cfg.momentum) optimizer.minimize(loss) fetch_list = fetch_list + [lr] for var in fetch_list: var.persistable = True gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0)) place = fluid.CUDAPlace(gpu_id) if cfg.use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) build_strategy = fluid.BuildStrategy() build_strategy.fuse_all_optimizer_ops = False build_strategy.fuse_elewise_add_act_ops = True exec_strategy = fluid.ExecutionStrategy() exec_strategy.num_iteration_per_drop_scope = 1 exe.run(startup_prog) if cfg.pretrained_model: checkpoint.load_and_fusebn(exe, train_prog, cfg.pretrained_model) compiled_train_prog = fluid.CompiledProgram(train_prog).with_data_parallel( loss_name=loss.name, build_strategy=build_strategy, exec_strategy=exec_strategy) shuffle = True shuffle_seed = None if num_trainers > 1: shuffle_seed = 1 if cfg.use_pyreader: train_reader = reader.train(batch_size=cfg.TRAIN.im_per_batch, total_batch_size=total_batch_size, padding_total=cfg.TRAIN.padding_minibatch, shuffle=shuffle, shuffle_seed=shuffle_seed) if num_trainers > 1: assert shuffle_seed is not None, \ "If num_trainers > 1, the shuffle_seed must be set, because " \ "the order of batch data generated by reader " \ "must be the same in the respective processes." # NOTE: the order of batch data generated by batch_reader # must be the same in the respective processes. if num_trainers > 1: train_reader = fluid.contrib.reader.distributed_batch_reader( train_reader) data_loader = model.data_loader data_loader.set_sample_list_generator(train_reader, places=place) else: if num_trainers > 1: shuffle = False train_reader = reader.train(batch_size=total_batch_size, shuffle=shuffle) feeder = fluid.DataFeeder(place=place, feed_list=model.feeds()) def train_loop(): data_loader.start() train_stats = TrainingStats(cfg.log_window, keys) try: start_time = time.time() prev_start_time = start_time for iter_id in range(cfg.max_iter): prev_start_time = start_time start_time = time.time() outs = exe.run(compiled_train_prog, fetch_list=[v.name for v in fetch_list]) stats = { k: np.array(v).mean() for k, v in zip(keys, outs[:-1]) } train_stats.update(stats) logs = train_stats.log() if iter_id % 10 == 0: strs = '{}, iter: {}, lr: {:.5f}, {}, time: {:.3f}'.format( now_time(), iter_id, np.mean(outs[-1]), logs, start_time - prev_start_time) print(strs) sys.stdout.flush() if (iter_id) % cfg.TRAIN.snapshot_iter == 0 and iter_id != 0: save_name = "{}".format(iter_id) checkpoint.save( exe, train_prog, os.path.join(cfg.model_save_dir, save_name)) if (iter_id) == cfg.max_iter: checkpoint.save( exe, train_prog, os.path.join(cfg.model_save_dir, "model_final")) break end_time = time.time() total_time = end_time - start_time last_loss = np.array(outs[0]).mean() except (StopIteration, fluid.core.EOFException): data_loader.reset() train_loop()