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()