import os import shutil import numpy as np import time import argparse import functools import paddle.fluid as fluid from pyramidbox import PyramidBox import reader from utility import add_arguments, print_arguments parser = argparse.ArgumentParser(description=__doc__) add_arg = functools.partial(add_arguments, argparser=parser) # yapf: disable add_arg('parallel', bool, True, "Whether use multi-GPU/threads or not.") add_arg('learning_rate', float, 0.001, "The start learning rate.") add_arg('batch_size', int, 16, "Minibatch size.") add_arg('num_passes', int, 160, "Epoch number.") add_arg('use_gpu', bool, True, "Whether use GPU.") add_arg('use_pyramidbox', bool, True, "Whether use PyramidBox model.") add_arg('model_save_dir', str, 'output', "The path to save model.") add_arg('resize_h', int, 640, "The resized image height.") add_arg('resize_w', int, 640, "The resized image width.") add_arg('with_mem_opt', bool, True, "Whether to use memory optimization or not.") add_arg('pretrained_model', str, './vgg_ilsvrc_16_fc_reduced/', "The init model path.") #yapf: enable def train(args, config, train_file_list, optimizer_method): learning_rate = args.learning_rate batch_size = args.batch_size num_passes = args.num_passes height = args.resize_h width = args.resize_w use_gpu = args.use_gpu use_pyramidbox = args.use_pyramidbox model_save_dir = args.model_save_dir pretrained_model = args.pretrained_model with_memory_optimization = args.with_mem_opt num_classes = 2 image_shape = [3, height, width] devices = os.getenv("CUDA_VISIBLE_DEVICES") or "" devices_num = len(devices.split(",")) fetches = [] network = PyramidBox(image_shape, num_classes, sub_network=use_pyramidbox) if use_pyramidbox: face_loss, head_loss, loss = network.train() fetches = [face_loss, head_loss] else: loss = network.vgg_ssd_loss() fetches = [loss] steps_per_pass = 12880 / batch_size boundaries = [steps_per_pass * 50, steps_per_pass * 80, steps_per_pass * 120, steps_per_pass * 140] values = [ learning_rate, learning_rate * 0.5, learning_rate * 0.25, learning_rate * 0.1, learning_rate * 0.01 ] if optimizer_method == "momentum": optimizer = fluid.optimizer.Momentum( learning_rate=fluid.layers.piecewise_decay( boundaries=boundaries, values=values), momentum=0.9, regularization=fluid.regularizer.L2Decay(0.0005), ) else: optimizer = fluid.optimizer.RMSProp( learning_rate=fluid.layers.piecewise_decay(boundaries, values), regularization=fluid.regularizer.L2Decay(0.0005), ) optimizer.minimize(loss) if with_memory_optimization: fluid.memory_optimize(fluid.default_main_program()) place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) exe.run(fluid.default_startup_program()) start_pass = 0 if pretrained_model: if pretrained_model.isdigit(): start_pass = int(pretrained_model) + 1 pretrained_model = os.path.join(model_save_dir, pretrained_model) print("Resume from %s " %(pretrained_model)) if not os.path.exists(pretrained_model): raise ValueError("The pre-trained model path [%s] does not exist." % (pretrained_model)) def if_exist(var): return os.path.exists(os.path.join(pretrained_model, var.name)) fluid.io.load_vars(exe, pretrained_model, predicate=if_exist) if args.parallel: train_exe = fluid.ParallelExecutor( use_cuda=use_gpu, loss_name=loss.name) train_reader = reader.train_batch_reader(config, train_file_list, batch_size=batch_size) def save_model(postfix): model_path = os.path.join(model_save_dir, postfix) if os.path.isdir(model_path): shutil.rmtree(model_path) print 'save models to %s' % (model_path) fluid.io.save_persistables(exe, model_path) def tensor(data, place, lod=None): t = fluid.core.LoDTensor() t.set(data, place) if lod: t.set_lod(lod) return t for pass_id in range(start_pass, num_passes): start_time = time.time() prev_start_time = start_time end_time = 0 for batch_id in range(steps_per_pass): im, face_box, head_box, labels, lod = next(train_reader) im_t = tensor(im, place) box1 = tensor(face_box, place, [lod]) box2 = tensor(head_box, place, [lod]) lbl_t = tensor(labels, place, [lod]) feeding = {'image': im_t, 'face_box': box1, 'head_box': box2, 'gt_label': lbl_t} prev_start_time = start_time start_time = time.time() if args.parallel: fetch_vars = train_exe.run(fetch_list=[v.name for v in fetches], feed=feeding) else: fetch_vars = exe.run(fluid.default_main_program(), feed=feeding, fetch_list=fetches) end_time = time.time() fetch_vars = [np.mean(np.array(v)) for v in fetch_vars] if batch_id % 10 == 0: if not args.use_pyramidbox: print("Pass {0}, batch {1}, loss {2}, time {3}".format( pass_id, batch_id, fetch_vars[0], start_time - prev_start_time)) else: print("Pass {0}, batch {1}, face loss {2}, head loss {3}, " \ "time {4}".format(pass_id, batch_id, fetch_vars[0], fetch_vars[1], start_time - prev_start_time)) if pass_id % 1 == 0 or pass_id == num_passes - 1: save_model(str(pass_id)) if __name__ == '__main__': args = parser.parse_args() print_arguments(args) data_dir = 'data/WIDER_train/images/' train_file_list = 'data/wider_face_split/wider_face_train_bbx_gt.txt' config = reader.Settings( data_dir=data_dir, resize_h=args.resize_h, resize_w=args.resize_w, apply_distort=True, apply_expand=False, mean_value=[104., 117., 123.], ap_version='11point') train(args, config, train_file_list, optimizer_method="momentum")