def main():
env = os.environ
cfg = load_config(FLAGS.config)
merge_config(FLAGS.opt)
check_config(cfg)
# check if set use_gpu=True in paddlepaddle cpu version
check_gpu(cfg.use_gpu)
main_arch = cfg.architecture
if cfg.use_gpu:
devices_num = fluid.core.get_cuda_device_count()
else:
devices_num = int(os.environ.get('CPU_NUM', 1))
if 'FLAGS_selected_gpus' in env:
device_id = int(env['FLAGS_selected_gpus'])
else:
device_id = 0
place = fluid.CUDAPlace(device_id) if cfg.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
# build program
model = create(main_arch)
inputs_def = cfg['TrainReader']['inputs_def']
train_feed_vars, train_loader = model.build_inputs(**inputs_def)
train_fetches = model.train(train_feed_vars)
loss = train_fetches['loss']
start_iter = 0
train_reader = create_reader(cfg.TrainReader,
(cfg.max_iters - start_iter) * devices_num,
cfg)
train_loader.set_sample_list_generator(train_reader, place)
# get all student variables
student_vars = []
for v in fluid.default_main_program().list_vars():
try:
student_vars.append((v.name, v.shape))
except:
pass
# uncomment the following lines to print all student variables
# print("="*50 + "student_model_vars" + "="*50)
# print(student_vars)
eval_prog = fluid.Program()
with fluid.program_guard(eval_prog, fluid.default_startup_program()):
with fluid.unique_name.guard():
model = create(main_arch)
inputs_def = cfg['EvalReader']['inputs_def']
test_feed_vars, eval_loader = model.build_inputs(**inputs_def)
fetches = model.eval(test_feed_vars)
eval_prog = eval_prog.clone(True)
eval_reader = create_reader(cfg.EvalReader)
eval_loader.set_sample_list_generator(eval_reader, place)
# parse eval fetches
extra_keys = []
if cfg.metric == 'COCO':
extra_keys = ['im_info', 'im_id', 'im_shape']
if cfg.metric == 'VOC':
extra_keys = ['gt_bbox', 'gt_class', 'is_difficult']
eval_keys, eval_values, eval_cls = parse_fetches(fetches, eval_prog,
extra_keys)
teacher_cfg = load_config(FLAGS.teacher_config)
merge_config(FLAGS.opt)
teacher_arch = teacher_cfg.architecture
teacher_program = fluid.Program()
teacher_startup_program = fluid.Program()
with fluid.program_guard(teacher_program, teacher_startup_program):
with fluid.unique_name.guard():
teacher_feed_vars = OrderedDict()
for name, var in train_feed_vars.items():
teacher_feed_vars[name] = teacher_program.global_block(
)._clone_variable(var, force_persistable=False)
model = create(teacher_arch)
train_fetches = model.train(teacher_feed_vars)
teacher_loss = train_fetches['loss']
# get all teacher variables
teacher_vars = []
for v in teacher_program.list_vars():
try:
teacher_vars.append((v.name, v.shape))
except:
pass
# uncomment the following lines to print all teacher variables
# print("="*50 + "teacher_model_vars" + "="*50)
# print(teacher_vars)
exe.run(teacher_startup_program)
assert FLAGS.teacher_pretrained, "teacher_pretrained should be set"
checkpoint.load_params(exe, teacher_program, FLAGS.teacher_pretrained)
teacher_program = teacher_program.clone(for_test=True)
cfg = load_config(FLAGS.config)
merge_config(FLAGS.opt)
data_name_map = {
'target0': 'target0',
'target1': 'target1',
'target2': 'target2',
'image': 'image',
'gt_bbox': 'gt_bbox',
'gt_class': 'gt_class',
'gt_score': 'gt_score'
}
merge(teacher_program, fluid.default_main_program(), data_name_map, place)
yolo_output_names = [
'strided_slice_0.tmp_0', 'strided_slice_1.tmp_0',
'strided_slice_2.tmp_0', 'strided_slice_3.tmp_0',
'strided_slice_4.tmp_0', 'transpose_0.tmp_0', 'strided_slice_5.tmp_0',
'strided_slice_6.tmp_0', 'strided_slice_7.tmp_0',
'strided_slice_8.tmp_0', 'strided_slice_9.tmp_0', 'transpose_2.tmp_0',
'strided_slice_10.tmp_0', 'strided_slice_11.tmp_0',
'strided_slice_12.tmp_0', 'strided_slice_13.tmp_0',
'strided_slice_14.tmp_0', 'transpose_4.tmp_0'
]
distill_pairs = [['teacher_conv2d_6.tmp_1', 'conv2d_20.tmp_1'],
['teacher_conv2d_14.tmp_1', 'conv2d_28.tmp_1'],
['teacher_conv2d_22.tmp_1', 'conv2d_36.tmp_1']]
distill_loss = l2_distill(
distill_pairs,
100) if not cfg.use_fine_grained_loss else split_distill(
yolo_output_names, 1000)
loss = distill_loss + loss
lr_builder = create('LearningRate')
optim_builder = create('OptimizerBuilder')
lr = lr_builder()
opt = optim_builder(lr)
opt.minimize(loss)
exe.run(fluid.default_startup_program())
fuse_bn = getattr(model.backbone, 'norm_type', None) == 'affine_channel'
ignore_params = cfg.finetune_exclude_pretrained_params \
if 'finetune_exclude_pretrained_params' in cfg else []
if FLAGS.resume_checkpoint:
checkpoint.load_checkpoint(exe, fluid.default_main_program(),
FLAGS.resume_checkpoint)
start_iter = checkpoint.global_step()
elif cfg.pretrain_weights and fuse_bn and not ignore_params:
checkpoint.load_and_fusebn(exe, fluid.default_main_program(),
cfg.pretrain_weights)
elif cfg.pretrain_weights:
checkpoint.load_params(exe,
fluid.default_main_program(),
cfg.pretrain_weights,
ignore_params=ignore_params)
build_strategy = fluid.BuildStrategy()
build_strategy.fuse_all_reduce_ops = False
build_strategy.fuse_all_optimizer_ops = False
# only enable sync_bn in multi GPU devices
sync_bn = getattr(model.backbone, 'norm_type', None) == 'sync_bn'
build_strategy.sync_batch_norm = sync_bn and devices_num > 1 \
and cfg.use_gpu
exec_strategy = fluid.ExecutionStrategy()
# iteration number when CompiledProgram tries to drop local execution scopes.
# Set it to be 1 to save memory usages, so that unused variables in
# local execution scopes can be deleted after each iteration.
exec_strategy.num_iteration_per_drop_scope = 1
parallel_main = fluid.CompiledProgram(
fluid.default_main_program()).with_data_parallel(
loss_name=loss.name,
build_strategy=build_strategy,
exec_strategy=exec_strategy)
compiled_eval_prog = fluid.compiler.CompiledProgram(eval_prog)
# whether output bbox is normalized in model output layer
is_bbox_normalized = False
if hasattr(model, 'is_bbox_normalized') and \
callable(model.is_bbox_normalized):
is_bbox_normalized = model.is_bbox_normalized()
map_type = cfg.map_type if 'map_type' in cfg else '11point'
best_box_ap_list = [0.0, 0] #[map, iter]
cfg_name = os.path.basename(FLAGS.config).split('.')[0]
save_dir = os.path.join(cfg.save_dir, cfg_name)
train_loader.start()
for step_id in range(start_iter, cfg.max_iters):
teacher_loss_np, distill_loss_np, loss_np, lr_np = exe.run(
parallel_main,
fetch_list=[
'teacher_' + teacher_loss.name, distill_loss.name, loss.name,
lr.name
])
if step_id % cfg.log_iter == 0:
logger.info(
"step {} lr {:.6f}, loss {:.6f}, distill_loss {:.6f}, teacher_loss {:.6f}"
.format(step_id, lr_np[0], loss_np[0], distill_loss_np[0],
teacher_loss_np[0]))
if step_id % cfg.snapshot_iter == 0 and step_id != 0 or step_id == cfg.max_iters - 1:
save_name = str(
step_id) if step_id != cfg.max_iters - 1 else "model_final"
checkpoint.save(exe, fluid.default_main_program(),
os.path.join(save_dir, save_name))
if FLAGS.save_inference:
feeded_var_names = ['image', 'im_size']
targets = list(fetches.values())
fluid.io.save_inference_model(save_dir + '/infer',
feeded_var_names, targets, exe,
eval_prog)
# eval
results = eval_run(exe, compiled_eval_prog, eval_loader, eval_keys,
eval_values, eval_cls, cfg)
resolution = None
box_ap_stats = eval_results(results, cfg.metric, cfg.num_classes,
resolution, is_bbox_normalized,
FLAGS.output_eval, map_type,
cfg['EvalReader']['dataset'])
if box_ap_stats[0] > best_box_ap_list[0]:
best_box_ap_list[0] = box_ap_stats[0]
best_box_ap_list[1] = step_id
checkpoint.save(exe, fluid.default_main_program(),
os.path.join(save_dir, "best_model"))
if FLAGS.save_inference:
feeded_var_names = ['image', 'im_size']
targets = list(fetches.values())
fluid.io.save_inference_model(save_dir + '/infer',
feeded_var_names, targets,
exe, eval_prog)
logger.info("Best test box ap: {}, in step: {}".format(
best_box_ap_list[0], best_box_ap_list[1]))
train_loader.reset()
def main():
env = os.environ
cfg = load_config(FLAGS.config)
merge_config(FLAGS.opt)
check_config(cfg)
# check if set use_gpu=True in paddlepaddle cpu version
check_gpu(cfg.use_gpu)
check_version()
main_arch = cfg.architecture
if cfg.use_gpu:
devices_num = fluid.core.get_cuda_device_count()
else:
devices_num = int(os.environ.get('CPU_NUM', 1))
if 'FLAGS_selected_gpus' in env:
device_id = int(env['FLAGS_selected_gpus'])
else:
device_id = 0
place = fluid.CUDAPlace(device_id) if cfg.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
# build program
model = create(main_arch)
inputs_def = cfg['TrainReader']['inputs_def']
train_feed_vars, train_loader = model.build_inputs(**inputs_def)
train_fetches = model.train(train_feed_vars)
loss = train_fetches['loss']
start_iter = 0
train_reader = create_reader(cfg.TrainReader,
(cfg.max_iters - start_iter) * devices_num,
cfg)
# When iterable mode, set set_sample_list_generator(train_reader, place)
train_loader.set_sample_list_generator(train_reader)
eval_prog = fluid.Program()
with fluid.program_guard(eval_prog, fluid.default_startup_program()):
with fluid.unique_name.guard():
model = create(main_arch)
inputs_def = cfg['EvalReader']['inputs_def']
test_feed_vars, eval_loader = model.build_inputs(**inputs_def)
fetches = model.eval(test_feed_vars)
eval_prog = eval_prog.clone(True)
eval_reader = create_reader(cfg.EvalReader)
# When iterable mode, set set_sample_list_generator(eval_reader, place)
eval_loader.set_sample_list_generator(eval_reader)
teacher_cfg = load_config(FLAGS.teacher_config)
merge_config(FLAGS.opt)
teacher_arch = teacher_cfg.architecture
teacher_program = fluid.Program()
teacher_startup_program = fluid.Program()
with fluid.program_guard(teacher_program, teacher_startup_program):
with fluid.unique_name.guard():
teacher_feed_vars = OrderedDict()
for name, var in train_feed_vars.items():
teacher_feed_vars[name] = teacher_program.global_block(
)._clone_variable(var, force_persistable=False)
model = create(teacher_arch)
train_fetches = model.train(teacher_feed_vars)
teacher_loss = train_fetches['loss']
exe.run(teacher_startup_program)
assert FLAGS.teacher_pretrained, "teacher_pretrained should be set"
checkpoint.load_params(exe, teacher_program, FLAGS.teacher_pretrained)
teacher_program = teacher_program.clone(for_test=True)
target_number = len(model.yolo_head.anchor_masks)
data_name_map = {
'image': 'image',
'gt_bbox': 'gt_bbox',
'gt_class': 'gt_class',
'gt_score': 'gt_score'
}
for i in range(target_number):
data_name_map['target{}'.format(i)] = 'target{}'.format(i)
merge(teacher_program, fluid.default_main_program(), data_name_map, place)
output_names = [
[
'strided_slice_0.tmp_0', 'strided_slice_1.tmp_0',
'strided_slice_2.tmp_0', 'strided_slice_3.tmp_0',
'strided_slice_4.tmp_0', 'transpose_0.tmp_0'
],
[
'strided_slice_5.tmp_0', 'strided_slice_6.tmp_0',
'strided_slice_7.tmp_0', 'strided_slice_8.tmp_0',
'strided_slice_9.tmp_0', 'transpose_2.tmp_0'
],
[
'strided_slice_10.tmp_0', 'strided_slice_11.tmp_0',
'strided_slice_12.tmp_0', 'strided_slice_13.tmp_0',
'strided_slice_14.tmp_0', 'transpose_4.tmp_0'
],
]
yolo_output_names = []
for i in range(target_number):
yolo_output_names.extend(output_names[i])
assert cfg.use_fine_grained_loss, \
"Only support use_fine_grained_loss=True, Please set it in config file or '-o use_fine_grained_loss=true'"
distill_loss = split_distill(yolo_output_names, 1000, target_number)
loss = distill_loss + loss
lr_builder = create('LearningRate')
optim_builder = create('OptimizerBuilder')
lr = lr_builder()
opt = optim_builder(lr)
opt.minimize(loss)
exe.run(fluid.default_startup_program())
checkpoint.load_params(exe, fluid.default_main_program(),
cfg.pretrain_weights)
assert FLAGS.pruned_params is not None, \
"FLAGS.pruned_params is empty!!! Please set it by '--pruned_params' option."
pruned_params = FLAGS.pruned_params.strip().split(",")
logger.info("pruned params: {}".format(pruned_params))
pruned_ratios = [float(n) for n in FLAGS.pruned_ratios.strip().split(",")]
logger.info("pruned ratios: {}".format(pruned_ratios))
assert len(pruned_params) == len(pruned_ratios), \
"The length of pruned params and pruned ratios should be equal."
assert pruned_ratios > [0] * len(pruned_ratios) and pruned_ratios < [1] * len(pruned_ratios), \
"The elements of pruned ratios should be in range (0, 1)."
assert FLAGS.prune_criterion in ['l1_norm', 'geometry_median'], \
"unsupported prune criterion {}".format(FLAGS.prune_criterion)
pruner = Pruner(criterion=FLAGS.prune_criterion)
distill_prog = pruner.prune(fluid.default_main_program(),
fluid.global_scope(),
params=pruned_params,
ratios=pruned_ratios,
place=place,
only_graph=False)[0]
base_flops = flops(eval_prog)
eval_prog = pruner.prune(eval_prog,
fluid.global_scope(),
params=pruned_params,
ratios=pruned_ratios,
place=place,
only_graph=True)[0]
pruned_flops = flops(eval_prog)
logger.info("FLOPs -{}; total FLOPs: {}; pruned FLOPs: {}".format(
float(base_flops - pruned_flops) / base_flops, base_flops,
pruned_flops))
build_strategy = fluid.BuildStrategy()
build_strategy.fuse_all_reduce_ops = False
build_strategy.fuse_all_optimizer_ops = False
build_strategy.fuse_elewise_add_act_ops = True
# only enable sync_bn in multi GPU devices
sync_bn = getattr(model.backbone, 'norm_type', None) == 'sync_bn'
build_strategy.sync_batch_norm = sync_bn and devices_num > 1 \
and cfg.use_gpu
exec_strategy = fluid.ExecutionStrategy()
# iteration number when CompiledProgram tries to drop local execution scopes.
# Set it to be 1 to save memory usages, so that unused variables in
# local execution scopes can be deleted after each iteration.
exec_strategy.num_iteration_per_drop_scope = 1
parallel_main = fluid.CompiledProgram(distill_prog).with_data_parallel(
loss_name=loss.name,
build_strategy=build_strategy,
exec_strategy=exec_strategy)
compiled_eval_prog = fluid.CompiledProgram(eval_prog)
# parse eval fetches
extra_keys = []
if cfg.metric == 'COCO':
extra_keys = ['im_info', 'im_id', 'im_shape']
if cfg.metric == 'VOC':
extra_keys = ['gt_bbox', 'gt_class', 'is_difficult']
eval_keys, eval_values, eval_cls = parse_fetches(fetches, eval_prog,
extra_keys)
# whether output bbox is normalized in model output layer
is_bbox_normalized = False
if hasattr(model, 'is_bbox_normalized') and \
callable(model.is_bbox_normalized):
is_bbox_normalized = model.is_bbox_normalized()
map_type = cfg.map_type if 'map_type' in cfg else '11point'
best_box_ap_list = [0.0, 0] #[map, iter]
cfg_name = os.path.basename(FLAGS.config).split('.')[0]
save_dir = os.path.join(cfg.save_dir, cfg_name)
train_loader.start()
for step_id in range(start_iter, cfg.max_iters):
teacher_loss_np, distill_loss_np, loss_np, lr_np = exe.run(
parallel_main,
fetch_list=[
'teacher_' + teacher_loss.name, distill_loss.name, loss.name,
lr.name
])
if step_id % cfg.log_iter == 0:
logger.info(
"step {} lr {:.6f}, loss {:.6f}, distill_loss {:.6f}, teacher_loss {:.6f}"
.format(step_id, lr_np[0], loss_np[0], distill_loss_np[0],
teacher_loss_np[0]))
if step_id % cfg.snapshot_iter == 0 and step_id != 0 or step_id == cfg.max_iters - 1:
save_name = str(
step_id) if step_id != cfg.max_iters - 1 else "model_final"
checkpoint.save(exe, distill_prog,
os.path.join(save_dir, save_name))
# eval
results = eval_run(exe, compiled_eval_prog, eval_loader, eval_keys,
eval_values, eval_cls, cfg)
resolution = None
box_ap_stats = eval_results(results, cfg.metric, cfg.num_classes,
resolution, is_bbox_normalized,
FLAGS.output_eval, map_type,
cfg['EvalReader']['dataset'])
if box_ap_stats[0] > best_box_ap_list[0]:
best_box_ap_list[0] = box_ap_stats[0]
best_box_ap_list[1] = step_id
checkpoint.save(exe, distill_prog,
os.path.join("./", "best_model"))
logger.info("Best test box ap: {}, in step: {}".format(
best_box_ap_list[0], best_box_ap_list[1]))
train_loader.reset()
def train(cfg):
# startup_prog = fluid.Program()
# train_prog = fluid.Program()
drop_last = True
dataset = SegDataset(
file_list=cfg.DATASET.TRAIN_FILE_LIST,
mode=ModelPhase.TRAIN,
shuffle=True,
data_dir=cfg.DATASET.DATA_DIR)
def data_generator():
if args.use_mpio:
data_gen = dataset.multiprocess_generator(
num_processes=cfg.DATALOADER.NUM_WORKERS,
max_queue_size=cfg.DATALOADER.BUF_SIZE)
else:
data_gen = dataset.generator()
batch_data = []
for b in data_gen:
batch_data.append(b)
if len(batch_data) == (cfg.BATCH_SIZE // cfg.NUM_TRAINERS):
for item in batch_data:
yield item[0], item[1], item[2]
batch_data = []
# If use sync batch norm strategy, drop last batch if number of samples
# in batch_data is less then cfg.BATCH_SIZE to avoid NCCL hang issues
if not cfg.TRAIN.SYNC_BATCH_NORM:
for item in batch_data:
yield item[0], item[1], item[2]
# Get device environment
# places = fluid.cuda_places() if args.use_gpu else fluid.cpu_places()
# place = places[0]
gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
place = fluid.CUDAPlace(gpu_id) if args.use_gpu else fluid.CPUPlace()
places = fluid.cuda_places() if args.use_gpu else fluid.cpu_places()
# Get number of GPU
dev_count = cfg.NUM_TRAINERS if cfg.NUM_TRAINERS > 1 else len(places)
print_info("#Device count: {}".format(dev_count))
# Make sure BATCH_SIZE can divided by GPU cards
assert cfg.BATCH_SIZE % dev_count == 0, (
'BATCH_SIZE:{} not divisble by number of GPUs:{}'.format(
cfg.BATCH_SIZE, dev_count))
# If use multi-gpu training mode, batch data will allocated to each GPU evenly
batch_size_per_dev = cfg.BATCH_SIZE // dev_count
print_info("batch_size_per_dev: {}".format(batch_size_per_dev))
data_loader, loss, lr, pred, grts, masks, image = build_model(
phase=ModelPhase.TRAIN)
data_loader.set_sample_generator(
data_generator, batch_size=batch_size_per_dev, drop_last=drop_last)
exe = fluid.Executor(place)
cfg.update_from_file(args.teacher_cfg_file)
# teacher_arch = teacher_cfg.architecture
teacher_program = fluid.Program()
teacher_startup_program = fluid.Program()
with fluid.program_guard(teacher_program, teacher_startup_program):
with fluid.unique_name.guard():
_, teacher_loss, _, _, _, _, _ = build_model(
teacher_program,
teacher_startup_program,
phase=ModelPhase.TRAIN,
image=image,
label=grts,
mask=masks)
exe.run(teacher_startup_program)
teacher_program = teacher_program.clone(for_test=True)
ckpt_dir = cfg.SLIM.KNOWLEDGE_DISTILL_TEACHER_MODEL_DIR
assert ckpt_dir is not None
print('load teacher model:', ckpt_dir)
if os.path.exists(ckpt_dir):
try:
fluid.load(teacher_program, os.path.join(ckpt_dir, 'model'), exe)
except:
fluid.io.load_params(exe, ckpt_dir, main_program=teacher_program)
# cfg = load_config(FLAGS.config)
cfg.update_from_file(args.cfg_file)
data_name_map = {
'image': 'image',
'label': 'label',
'mask': 'mask',
}
merge(teacher_program, fluid.default_main_program(), data_name_map, place)
distill_pairs = [[
'teacher_bilinear_interp_2.tmp_0', 'bilinear_interp_0.tmp_0'
]]
def distill(pairs, weight):
"""
Add 3 pairs of distillation losses, each pair of feature maps is the
input of teacher and student's yolov3_loss respectively
"""
loss = l2_loss(pairs[0][0], pairs[0][1])
weighted_loss = loss * weight
return weighted_loss
distill_loss = distill(distill_pairs, 0.1)
cfg.update_from_file(args.cfg_file)
optimizer = solver.Solver(None, None)
all_loss = loss + distill_loss
lr = optimizer.optimise(all_loss)
exe.run(fluid.default_startup_program())
exec_strategy = fluid.ExecutionStrategy()
# Clear temporary variables every 100 iteration
if args.use_gpu:
exec_strategy.num_threads = fluid.core.get_cuda_device_count()
exec_strategy.num_iteration_per_drop_scope = 100
build_strategy = fluid.BuildStrategy()
build_strategy.fuse_all_reduce_ops = False
build_strategy.fuse_all_optimizer_ops = False
build_strategy.fuse_elewise_add_act_ops = True
if cfg.NUM_TRAINERS > 1 and args.use_gpu:
dist_utils.prepare_for_multi_process(exe, build_strategy,
fluid.default_main_program())
exec_strategy.num_threads = 1
if cfg.TRAIN.SYNC_BATCH_NORM and args.use_gpu:
if dev_count > 1:
# Apply sync batch norm strategy
print_info("Sync BatchNorm strategy is effective.")
build_strategy.sync_batch_norm = True
else:
print_info(
"Sync BatchNorm strategy will not be effective if GPU device"
" count <= 1")
compiled_train_prog = fluid.CompiledProgram(
fluid.default_main_program()).with_data_parallel(
loss_name=all_loss.name,
exec_strategy=exec_strategy,
build_strategy=build_strategy)
# Resume training
begin_epoch = cfg.SOLVER.BEGIN_EPOCH
if cfg.TRAIN.RESUME_MODEL_DIR:
begin_epoch = load_checkpoint(exe, fluid.default_main_program())
# Load pretrained model
elif os.path.exists(cfg.TRAIN.PRETRAINED_MODEL_DIR):
load_pretrained_weights(exe, fluid.default_main_program(),
cfg.TRAIN.PRETRAINED_MODEL_DIR)
else:
print_info(
'Pretrained model dir {} not exists, training from scratch...'.
format(cfg.TRAIN.PRETRAINED_MODEL_DIR))
#fetch_list = [avg_loss.name, lr.name]
fetch_list = [
loss.name, 'teacher_' + teacher_loss.name, distill_loss.name, lr.name
]
if args.debug:
# Fetch more variable info and use streaming confusion matrix to
# calculate IoU results if in debug mode
np.set_printoptions(
precision=4, suppress=True, linewidth=160, floatmode="fixed")
fetch_list.extend([pred.name, grts.name, masks.name])
cm = ConfusionMatrix(cfg.DATASET.NUM_CLASSES, streaming=True)
if args.use_vdl:
if not args.vdl_log_dir:
print_info("Please specify the log directory by --vdl_log_dir.")
exit(1)
from visualdl import LogWriter
log_writer = LogWriter(args.vdl_log_dir)
# trainer_id = int(os.getenv("PADDLE_TRAINER_ID", 0))
# num_trainers = int(os.environ.get('PADDLE_TRAINERS_NUM', 1))
step = 0
all_step = cfg.DATASET.TRAIN_TOTAL_IMAGES // cfg.BATCH_SIZE
if cfg.DATASET.TRAIN_TOTAL_IMAGES % cfg.BATCH_SIZE and drop_last != True:
all_step += 1
all_step *= (cfg.SOLVER.NUM_EPOCHS - begin_epoch + 1)
avg_loss = 0.0
avg_t_loss = 0.0
avg_d_loss = 0.0
best_mIoU = 0.0
timer = Timer()
timer.start()
if begin_epoch > cfg.SOLVER.NUM_EPOCHS:
raise ValueError(
("begin epoch[{}] is larger than cfg.SOLVER.NUM_EPOCHS[{}]").format(
begin_epoch, cfg.SOLVER.NUM_EPOCHS))
if args.use_mpio:
print_info("Use multiprocess reader")
else:
print_info("Use multi-thread reader")
for epoch in range(begin_epoch, cfg.SOLVER.NUM_EPOCHS + 1):
data_loader.start()
while True:
try:
if args.debug:
# Print category IoU and accuracy to check whether the
# traning process is corresponed to expectation
loss, lr, pred, grts, masks = exe.run(
program=compiled_train_prog,
fetch_list=fetch_list,
return_numpy=True)
cm.calculate(pred, grts, masks)
avg_loss += np.mean(np.array(loss))
step += 1
if step % args.log_steps == 0:
speed = args.log_steps / timer.elapsed_time()
avg_loss /= args.log_steps
category_acc, mean_acc = cm.accuracy()
category_iou, mean_iou = cm.mean_iou()
print_info((
"epoch={} step={} lr={:.5f} loss={:.4f} acc={:.5f} mIoU={:.5f} step/sec={:.3f} | ETA {}"
).format(epoch, step, lr[0], avg_loss, mean_acc,
mean_iou, speed,
calculate_eta(all_step - step, speed)))
print_info("Category IoU: ", category_iou)
print_info("Category Acc: ", category_acc)
if args.use_vdl:
log_writer.add_scalar('Train/mean_iou', mean_iou,
step)
log_writer.add_scalar('Train/mean_acc', mean_acc,
step)
log_writer.add_scalar('Train/loss', avg_loss, step)
log_writer.add_scalar('Train/lr', lr[0], step)
log_writer.add_scalar('Train/step/sec', speed, step)
sys.stdout.flush()
avg_loss = 0.0
cm.zero_matrix()
timer.restart()
else:
# If not in debug mode, avoid unnessary log and calculate
loss, t_loss, d_loss, lr = exe.run(
program=compiled_train_prog,
fetch_list=fetch_list,
return_numpy=True)
avg_loss += np.mean(np.array(loss))
avg_t_loss += np.mean(np.array(t_loss))
avg_d_loss += np.mean(np.array(d_loss))
step += 1
if step % args.log_steps == 0 and cfg.TRAINER_ID == 0:
avg_loss /= args.log_steps
avg_t_loss /= args.log_steps
avg_d_loss /= args.log_steps
speed = args.log_steps / timer.elapsed_time()
print((
"epoch={} step={} lr={:.5f} loss={:.4f} teacher loss={:.4f} distill loss={:.4f} step/sec={:.3f} | ETA {}"
).format(epoch, step, lr[0], avg_loss, avg_t_loss,
avg_d_loss, speed,
calculate_eta(all_step - step, speed)))
if args.use_vdl:
log_writer.add_scalar('Train/loss', avg_loss, step)
log_writer.add_scalar('Train/lr', lr[0], step)
log_writer.add_scalar('Train/speed', speed, step)
sys.stdout.flush()
avg_loss = 0.0
avg_t_loss = 0.0
avg_d_loss = 0.0
timer.restart()
except fluid.core.EOFException:
data_loader.reset()
break
except Exception as e:
print(e)
if (epoch % cfg.TRAIN.SNAPSHOT_EPOCH == 0
or epoch == cfg.SOLVER.NUM_EPOCHS) and cfg.TRAINER_ID == 0:
ckpt_dir = save_checkpoint(fluid.default_main_program(), epoch)
if args.do_eval:
print("Evaluation start")
_, mean_iou, _, mean_acc = evaluate(
cfg=cfg,
ckpt_dir=ckpt_dir,
use_gpu=args.use_gpu,
use_mpio=args.use_mpio)
if args.use_vdl:
log_writer.add_scalar('Evaluate/mean_iou', mean_iou, step)
log_writer.add_scalar('Evaluate/mean_acc', mean_acc, step)
if mean_iou > best_mIoU:
best_mIoU = mean_iou
update_best_model(ckpt_dir)
print_info("Save best model {} to {}, mIoU = {:.4f}".format(
ckpt_dir,
os.path.join(cfg.TRAIN.MODEL_SAVE_DIR, 'best_model'),
mean_iou))
# Use VisualDL to visualize results
if args.use_vdl and cfg.DATASET.VIS_FILE_LIST is not None:
visualize(
cfg=cfg,
use_gpu=args.use_gpu,
vis_file_list=cfg.DATASET.VIS_FILE_LIST,
vis_dir="visual",
ckpt_dir=ckpt_dir,
log_writer=log_writer)
if cfg.TRAINER_ID == 0:
ckpt_dir = save_checkpoint(fluid.default_main_program(), epoch)
# save final model
if cfg.TRAINER_ID == 0:
save_checkpoint(fluid.default_main_program(), 'final')
