# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import time import numpy as np import argparse import functools import shutil import math import multiprocessing from paddle.fluid.contrib.slim import Compressor import paddle import paddle.fluid as fluid import reader from mobilenet_ssd import build_mobilenet_ssd from utility import add_arguments, print_arguments, check_cuda parser = argparse.ArgumentParser(description=__doc__) add_arg = functools.partial(add_arguments, argparser=parser) # yapf: disable add_arg('learning_rate', float, 0.001, "Learning rate.") add_arg('batch_size', int, 64, "Minibatch size of all devices.") add_arg('epoc_num', int, 120, "Epoch number.") add_arg('use_gpu', bool, True, "Whether use GPU.") add_arg('parallel', bool, True, "Whether train in parallel on multi-devices.") add_arg('dataset', str, 'pascalvoc', "dataset can be coco2014, coco2017, and pascalvoc.") add_arg('model_save_dir', str, 'model', "The path to save model.") add_arg('pretrained_model', str, 'pretrained/ssd_mobilenet_v1_coco/', "The init model path.") add_arg('ap_version', str, '11point', "mAP version can be integral or 11point.") add_arg('image_shape', str, '3,300,300', "Input image shape.") add_arg('mean_BGR', str, '127.5,127.5,127.5', "Mean value for B,G,R channel which will be subtracted.") add_arg('data_dir', str, 'data/pascalvoc', "Data directory.") add_arg('use_multiprocess', bool, True, "Whether use multi-process for data preprocessing.") add_arg('enable_ce', bool, False, "Whether use CE to evaluate the model.") #yapf: enable train_parameters = { "pascalvoc": { "train_images": 16551, "image_shape": [3, 300, 300], "class_num": 21, "batch_size": 64, "lr": 0.001, "lr_epochs": [40, 60, 80, 100], "lr_decay": [1, 0.5, 0.25, 0.1, 0.01], "ap_version": '11point', }, "coco2014": { "train_images": 82783, "image_shape": [3, 300, 300], "class_num": 91, "batch_size": 64, "lr": 0.001, "lr_epochs": [12, 19], "lr_decay": [1, 0.5, 0.25], "ap_version": 'integral', # should use eval_coco_map.py to test model }, "coco2017": { "train_images": 118287, "image_shape": [3, 300, 300], "class_num": 91, "batch_size": 64, "lr": 0.001, "lr_epochs": [12, 19], "lr_decay": [1, 0.5, 0.25], "ap_version": 'integral', # should use eval_coco_map.py to test model } } def optimizer_setting(train_params): batch_size = train_params["batch_size"] iters = train_params["train_images"] // batch_size lr = train_params["lr"] boundaries = [i * iters for i in train_params["lr_epochs"]] values = [ i * lr for i in train_params["lr_decay"]] optimizer = fluid.optimizer.RMSProp( learning_rate=fluid.layers.piecewise_decay(boundaries, values), regularization=fluid.regularizer.L2Decay(0.00005), ) return optimizer def build_program(main_prog, startup_prog, train_params, is_train): image_shape = train_params['image_shape'] class_num = train_params['class_num'] ap_version = train_params['ap_version'] outs = [] with fluid.program_guard(main_prog, startup_prog): py_reader = fluid.layers.py_reader( capacity=64, shapes=[[-1] + image_shape, [-1, 4], [-1, 1], [-1, 1]], lod_levels=[0, 1, 1, 1], dtypes=["float32", "float32", "int32", "int32"], use_double_buffer=True) with fluid.unique_name.guard(): image, gt_box, gt_label, difficult = fluid.layers.read_file(py_reader) locs, confs, box, box_var = build_mobilenet_ssd(image, class_num, image_shape) if is_train: with fluid.unique_name.guard("train"): loss = fluid.layers.ssd_loss(locs, confs, gt_box, gt_label, box, box_var) loss = fluid.layers.reduce_sum(loss) optimizer = optimizer_setting(train_params) optimizer.minimize(loss) outs = [py_reader, loss] else: with fluid.unique_name.guard("inference"): nmsed_out = fluid.layers.detection_output( locs, confs, box, box_var, nms_threshold=0.45) gt_label = fluid.layers.cast(x=gt_label, dtype=gt_box.dtype) if difficult: difficult = fluid.layers.cast(x=difficult, dtype=gt_box.dtype) gt_label = fluid.layers.reshape(gt_label, [-1, 1]) difficult = fluid.layers.reshape(difficult, [-1, 1]) label = fluid.layers.concat([gt_label, difficult, gt_box], axis=1) else: label = fluid.layers.concat([gt_label, gt_box], axis=1) map_var = fluid.layers.detection.detection_map( nmsed_out, label, class_num, background_label=0, overlap_threshold=0.5, evaluate_difficult=False, ap_version=ap_version) # nmsed_out and image is used to save mode for inference outs = [py_reader, map_var, nmsed_out, image] return outs def train(args, data_args, train_params, train_file_list, val_file_list): model_save_dir = args.model_save_dir pretrained_model = args.pretrained_model use_gpu = args.use_gpu parallel = args.parallel enable_ce = args.enable_ce is_shuffle = True if not use_gpu: devices_num = int(os.environ.get('CPU_NUM', multiprocessing.cpu_count())) else: devices_num = fluid.core.get_cuda_device_count() batch_size = train_params['batch_size'] epoc_num = train_params['epoc_num'] batch_size_per_device = batch_size // devices_num num_workers = 8 startup_prog = fluid.Program() train_prog = fluid.Program() test_prog = fluid.Program() train_py_reader, loss = build_program( main_prog=train_prog, startup_prog=startup_prog, train_params=train_params, is_train=True) test_py_reader, map_var, _, _ = build_program( main_prog=test_prog, startup_prog=startup_prog, train_params=train_params, is_train=False) test_prog = test_prog.clone(for_test=True) for param in train_prog.global_block().all_parameters(): if 'conv' in param.name: print param.name, param.shape place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) exe.run(startup_prog) if pretrained_model: def if_exist(var): return os.path.exists(os.path.join(pretrained_model, var.name)) fluid.io.load_vars(exe, pretrained_model, main_program=train_prog, predicate=if_exist) test_reader = reader.test(data_args, val_file_list, batch_size) test_py_reader.decorate_paddle_reader(test_reader) train_reader = reader.train(data_args, train_file_list, batch_size_per_device, shuffle=is_shuffle, use_multiprocess=args.use_multiprocess, num_workers=num_workers, enable_ce=enable_ce) train_py_reader.decorate_paddle_reader(train_reader) train_fetch_list=[("loss", loss.name)] val_fetch_list=[("map", map_var.name)] compressor = Compressor( place, fluid.global_scope(), train_prog, train_reader=train_py_reader, train_feed_list=None, train_fetch_list=train_fetch_list, eval_program=test_prog, eval_reader=test_py_reader, eval_feed_list=None, eval_fetch_list=val_fetch_list, train_optimizer=None) compressor.config('./compress.yaml') compressor.run() def main(): args = parser.parse_args() print_arguments(args) check_cuda(args.use_gpu) data_dir = args.data_dir dataset = args.dataset assert dataset in ['pascalvoc', 'coco2014', 'coco2017'] # for pascalvoc label_file = 'label_list' train_file_list = 'trainval.txt' val_file_list = 'test.txt' if dataset == 'coco2014': train_file_list = 'annotations/instances_train2014.json' val_file_list = 'annotations/instances_val2014.json' elif dataset == 'coco2017': train_file_list = 'annotations/instances_train2017.json' val_file_list = 'annotations/instances_val2017.json' mean_BGR = [float(m) for m in args.mean_BGR.split(",")] image_shape = [int(m) for m in args.image_shape.split(",")] train_parameters[dataset]['image_shape'] = image_shape train_parameters[dataset]['batch_size'] = args.batch_size train_parameters[dataset]['lr'] = args.learning_rate train_parameters[dataset]['epoc_num'] = args.epoc_num train_parameters[dataset]['ap_version'] = args.ap_version data_args = reader.Settings( dataset=args.dataset, data_dir=data_dir, label_file=label_file, resize_h=image_shape[1], resize_w=image_shape[2], mean_value=mean_BGR, apply_distort=True, apply_expand=True, ap_version = args.ap_version) train(args, data_args, train_parameters[dataset], train_file_list=train_file_list, val_file_list=val_file_list) if __name__ == '__main__': main()