From 385b33bdb408b347f20168dd2d58aaae1bad1baa Mon Sep 17 00:00:00 2001 From: jerrywgz Date: Thu, 20 Sep 2018 20:48:57 +0800 Subject: [PATCH] add infer (#1274) * refine test reader return * add infer and parse_args --- fluid/faster_rcnn/eval_coco_map.py | 151 ++++------------------ fluid/faster_rcnn/eval_helper.py | 55 ++++++++ fluid/faster_rcnn/models/model_builder.py | 11 +- fluid/faster_rcnn/reader.py | 32 +++-- fluid/faster_rcnn/roidbs.py | 3 +- fluid/faster_rcnn/train.py | 92 +++++-------- fluid/faster_rcnn/utility.py | 50 +++++++ 7 files changed, 197 insertions(+), 197 deletions(-) diff --git a/fluid/faster_rcnn/eval_coco_map.py b/fluid/faster_rcnn/eval_coco_map.py index 58de4426..126e6cdc 100644 --- a/fluid/faster_rcnn/eval_coco_map.py +++ b/fluid/faster_rcnn/eval_coco_map.py @@ -4,57 +4,31 @@ import numpy as np import argparse import functools from eval_helper import get_nmsed_box +from eval_helper import get_dt_res import paddle import paddle.fluid as fluid import reader -from utility import add_arguments, print_arguments -from PIL import Image -from PIL import ImageDraw -from PIL import ImageFont +from utility import print_arguments, parse_args # A special mAP metric for COCO dataset, which averages AP in different IoUs. -# To use this eval_cocoMAP.py, [cocoapi](https://github.com/cocodataset/cocoapi) is needed. +# To use this eval_coco_map.py, [cocoapi](https://github.com/cocodataset/cocoapi) is needed. import models.model_builder as model_builder import models.resnet as resnet import json from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval, Params -parser = argparse.ArgumentParser(description=__doc__) -add_arg = functools.partial(add_arguments, argparser=parser) -# yapf: disable -add_arg('dataset', str, 'coco2017', "coco2014, coco2017.") -add_arg('batch_size', int, 1, "Minibatch size.") -add_arg('use_gpu', bool, True, "Whether use GPU.") -add_arg('data_dir', str, 'data/COCO17', "The data root path.") -add_arg('model_dir', str, '', "The model path.") -add_arg('nms_threshold', float, 0.5, "NMS threshold.") -add_arg('score_threshold', float, 0.05, "score threshold for NMS.") -add_arg('confs_threshold', float, 9., "Confidence threshold to draw bbox.") -add_arg('image_path', str, '', "The image used to inference and visualize.") -add_arg('anchor_sizes', int, [32,64,128,256,512], "The size of anchors.") -add_arg('aspect_ratios', float, [0.5,1.0,2.0], "The ratio of anchors.") -add_arg('ap_version', str, 'cocoMAP', "cocoMAP.") -add_arg('max_size', int, 1333, "The resized image height.") -add_arg('scales', int, [800], "The resized image height.") -add_arg('mean_value', float, [102.9801, 115.9465, 122.7717], "pixel mean") -add_arg('class_num', int, 81, "Class number.") -add_arg('variance', float, [1.,1.,1.,1.], "The variance of anchors.") -# yapf: enable +def eval(cfg): - -def eval(args): - - if '2014' in args.dataset: + if '2014' in cfg.dataset: test_list = 'annotations/instances_val2014.json' - elif '2017' in args.dataset: + elif '2017' in cfg.dataset: test_list = 'annotations/instances_val2017.json' - image_shape = [3, args.max_size, args.max_size] - class_nums = args.class_num - batch_size = args.batch_size - - cocoGt = COCO(os.path.join(data_args.data_dir, test_list)) + image_shape = [3, cfg.max_size, cfg.max_size] + class_nums = cfg.class_num + batch_size = cfg.batch_size + cocoGt = COCO(os.path.join(cfg.data_dir, test_list)) numId_to_catId_map = {i + 1: v for i, v in enumerate(cocoGt.getCatIds())} category_ids = cocoGt.getCatIds() label_list = { @@ -62,61 +36,42 @@ def eval(args): for item in cocoGt.loadCats(category_ids) } label_list[0] = ['background'] - print(label_list) model = model_builder.FasterRCNN( - cfg=args, + cfg=cfg, add_conv_body_func=resnet.add_ResNet50_conv4_body, add_roi_box_head_func=resnet.add_ResNet_roi_conv5_head, use_pyreader=False, - is_train=False, - use_random=False) + is_train=False) model.build_model(image_shape) - rpn_rois, scores, locs = model.eval_out() - confs = fluid.layers.softmax(scores, use_cudnn=False) - place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace() + rpn_rois, confs, locs = model.eval_out() + place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) # yapf: disable - if args.model_dir: + if cfg.pretrained_model: def if_exist(var): - return os.path.exists(os.path.join(args.model_dir, var.name)) - fluid.io.load_vars(exe, args.model_dir, predicate=if_exist) + return os.path.exists(os.path.join(cfg.pretrained_model, var.name)) + fluid.io.load_vars(exe, cfg.pretrained_model, predicate=if_exist) # yapf: enable - test_reader = reader.test(args, batch_size) + test_reader = reader.test(cfg, batch_size) feeder = fluid.DataFeeder(place=place, feed_list=model.feeds()) dts_res = [] fetch_list = [rpn_rois, confs, locs] - for batch_id, data in enumerate(test_reader()): + for batch_id, batch_data in enumerate(test_reader()): start = time.time() - #image, gt_box, gt_label, is_crowd, im_info, im_id = data[0] im_info = [] - image = [] - for i in range(len(data)): - im_info.append(data[i][4]) - image.append(data[i][0]) - image_t = fluid.core.LoDTensor() - image_t.set(image, place) - - im_info_t = fluid.core.LoDTensor() - im_info_t.set(im_info, place) - - feeding = {} - feeding['image'] = image_t - feeding['im_info'] = im_info_t + for data in batch_data: + im_info.append(data[1]) rpn_rois_v, confs_v, locs_v = exe.run( fetch_list=[v.name for v in fetch_list], - #feed=feeder.feed(data), - feed=feeding, + feed=feeder.feed(batch_data), return_numpy=False) - new_lod, nmsed_out = get_nmsed_box(args, rpn_rois_v, confs_v, locs_v, + new_lod, nmsed_out = get_nmsed_box(cfg, rpn_rois_v, confs_v, locs_v, class_nums, im_info, numId_to_catId_map) - for i in range(len(data)): - if str(data[i][5]) in args.image_path: - draw_bounding_box_on_image(args.image_path, nmsed_out, - args.confs_threshold, label_list) - dts_res += get_dt_res(new_lod, nmsed_out, data) + + dts_res += get_dt_res(batch_size, new_lod, nmsed_out, batch_data) end = time.time() print('batch id: {}, time: {}'.format(batch_id, end - start)) with open("detection_result.json", 'w') as outfile: @@ -124,67 +79,13 @@ def eval(args): print("start evaluate using coco api") cocoDt = cocoGt.loadRes("detection_result.json") cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') - #cocoEval.params.imgIds = im_id cocoEval.evaluate() cocoEval.accumulate() cocoEval.summarize() -def get_dt_res(lod, nmsed_out, data): - dts_res = [] - nmsed_out_v = np.array(nmsed_out) - assert (len(lod) == args.batch_size + 1), \ - "Error Lod Tensor offset dimension. Lod({}) vs. batch_size({})"\ - .format(len(lod), batch_size) - k = 0 - for i in range(args.batch_size): - dt_num_this_img = lod[i + 1] - lod[i] - image_id = int(data[i][-1]) - image_width = int(data[i][4][1]) - image_height = int(data[i][4][2]) - for j in range(dt_num_this_img): - dt = nmsed_out_v[k] - k = k + 1 - xmin, ymin, xmax, ymax, score, category_id = dt.tolist() - w = xmax - xmin + 1 - h = ymax - ymin + 1 - bbox = [xmin, ymin, w, h] - dt_res = { - 'image_id': image_id, - 'category_id': category_id, - 'bbox': bbox, - 'score': score - } - dts_res.append(dt_res) - return dts_res - - -def draw_bounding_box_on_image(image_path, nms_out, confs_threshold, - label_list): - image = Image.open(image_path) - draw = ImageDraw.Draw(image) - im_width, im_height = image.size - - for dt in nms_out: - xmin, ymin, xmax, ymax, score, category_id = dt.tolist() - if score < confs_threshold: - continue - bbox = dt[:4] - xmin, ymin, xmax, ymax = bbox - draw.line( - [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), - (xmin, ymin)], - width=4, - fill='red') - if image.mode == 'RGB': - draw.text((xmin, ymin), label_list[int(category_id)], (255, 255, 0)) - image_name = image_path.split('/')[-1] - print("image with bbox drawed saved as {}".format(image_name)) - image.save(image_name) - - if __name__ == '__main__': - args = parser.parse_args() + args = parse_args() print_arguments(args) data_args = reader.Settings(args) diff --git a/fluid/faster_rcnn/eval_helper.py b/fluid/faster_rcnn/eval_helper.py index 4bc1ea93..bb25105a 100644 --- a/fluid/faster_rcnn/eval_helper.py +++ b/fluid/faster_rcnn/eval_helper.py @@ -3,6 +3,9 @@ import numpy as np import paddle.fluid as fluid import math import box_utils +from PIL import Image +from PIL import ImageDraw +from PIL import ImageFont def box_decoder(target_box, prior_box, prior_box_var): @@ -102,3 +105,55 @@ def get_nmsed_box(args, rpn_rois, confs, locs, class_nums, im_info, labels = im_results_n[:, -1] im_results = np.vstack([im_results[k] for k in range(len(lod) - 1)]) return new_lod, im_results + + +def get_dt_res(batch_size, lod, nmsed_out, data): + dts_res = [] + nmsed_out_v = np.array(nmsed_out) + assert (len(lod) == batch_size + 1), \ + "Error Lod Tensor offset dimension. Lod({}) vs. batch_size({})"\ + .format(len(lod), batch_size) + k = 0 + for i in range(batch_size): + dt_num_this_img = lod[i + 1] - lod[i] + image_id = int(data[i][-1]) + image_width = int(data[i][1][1]) + image_height = int(data[i][1][2]) + for j in range(dt_num_this_img): + dt = nmsed_out_v[k] + k = k + 1 + xmin, ymin, xmax, ymax, score, category_id = dt.tolist() + w = xmax - xmin + 1 + h = ymax - ymin + 1 + bbox = [xmin, ymin, w, h] + dt_res = { + 'image_id': image_id, + 'category_id': category_id, + 'bbox': bbox, + 'score': score + } + dts_res.append(dt_res) + return dts_res + + +def draw_bounding_box_on_image(image_path, nms_out, draw_threshold, label_list): + image = Image.open(image_path) + draw = ImageDraw.Draw(image) + im_width, im_height = image.size + + for dt in nms_out: + xmin, ymin, xmax, ymax, score, category_id = dt.tolist() + if score < draw_threshold: + continue + bbox = dt[:4] + xmin, ymin, xmax, ymax = bbox + draw.line( + [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), + (xmin, ymin)], + width=4, + fill='red') + if image.mode == 'RGB': + draw.text((xmin, ymin), label_list[int(category_id)], (255, 255, 0)) + image_name = image_path.split('/')[-1] + print("image with bbox drawed saved as {}".format(image_name)) + image.save(image_name) diff --git a/fluid/faster_rcnn/models/model_builder.py b/fluid/faster_rcnn/models/model_builder.py index de261bcc..80772907 100644 --- a/fluid/faster_rcnn/models/model_builder.py +++ b/fluid/faster_rcnn/models/model_builder.py @@ -37,7 +37,8 @@ class FasterRCNN(object): return loss_cls, loss_bbox, rpn_cls_loss, rpn_reg_loss, def eval_out(self): - return [self.rpn_rois, self.cls_score, self.bbox_pred] + cls_prob = fluid.layers.softmax(self.cls_score, use_cudnn=False) + return [self.rpn_rois, cls_prob, self.bbox_pred] def build_input(self, image_shape): if self.use_pyreader: @@ -71,6 +72,8 @@ class FasterRCNN(object): name='im_id', shape=[1], dtype='int32') def feeds(self): + if not self.is_train: + return [self.image, self.im_info, self.im_id] return [ self.image, self.gt_box, self.gt_label, self.is_crowd, self.im_info, self.im_id @@ -156,7 +159,7 @@ class FasterRCNN(object): is_crowd=self.is_crowd, gt_boxes=self.gt_box, im_info=self.im_info, - batch_size_per_im=512, + batch_size_per_im=self.cfg.batch_size_per_im, fg_fraction=0.25, fg_thresh=0.5, bg_thresh_hi=0.5, @@ -215,8 +218,8 @@ class FasterRCNN(object): # logits=cls_score, # label=labels_int64 # ) - softmax = fluid.layers.softmax(self.cls_score, use_cudnn=False) - loss_cls = fluid.layers.cross_entropy(softmax, labels_int64) + cls_prob = fluid.layers.softmax(self.cls_score, use_cudnn=False) + loss_cls = fluid.layers.cross_entropy(cls_prob, labels_int64) loss_cls = fluid.layers.reduce_mean(loss_cls) loss_bbox = fluid.layers.smooth_l1( x=self.bbox_pred, diff --git a/fluid/faster_rcnn/reader.py b/fluid/faster_rcnn/reader.py index 04bc0e7a..f134230c 100644 --- a/fluid/faster_rcnn/reader.py +++ b/fluid/faster_rcnn/reader.py @@ -59,13 +59,14 @@ def coco(settings, padding_total=False, shuffle=False): total_batch_size = total_batch_size if total_batch_size else batch_size - assert total_batch_size % batch_size == 0 + if mode != 'infer': + assert total_batch_size % batch_size == 0 if mode == 'train': settings.train_file_list = os.path.join(settings.data_dir, settings.train_file_list) settings.train_data_dir = os.path.join(settings.data_dir, settings.train_data_dir) - elif mode == 'test': + elif mode == 'test' or mode == 'infer': settings.val_file_list = os.path.join(settings.data_dir, settings.val_file_list) settings.val_data_dir = os.path.join(settings.data_dir, @@ -75,12 +76,14 @@ def coco(settings, print("{} on {} with {} roidbs".format(mode, settings.dataset, len(roidbs))) - def roidb_reader(roidb): + def roidb_reader(roidb, mode): im, im_scales = data_utils.get_image_blob(roidb, settings) im_id = roidb['id'] im_height = np.round(roidb['height'] * im_scales) im_width = np.round(roidb['width'] * im_scales) im_info = np.array([im_height, im_width, im_scales], dtype=np.float32) + if mode == 'test' or mode == 'infer': + return im, im_info, im_id gt_boxes = roidb['gt_boxes'].astype('float32') gt_classes = roidb['gt_classes'].astype('int32') is_crowd = roidb['is_crowd'].astype('int32') @@ -113,7 +116,7 @@ def coco(settings, if roidb_cur >= len(roidbs): roidb_perm = deque(np.random.permutation(roidbs)) im, gt_boxes, gt_classes, is_crowd, im_info, im_id = roidb_reader( - roidb) + roidb, mode) if gt_boxes.shape[0] == 0: continue batch_out.append( @@ -133,16 +136,23 @@ def coco(settings, yield sub_batch_out sub_batch_out = [] batch_out = [] - else: + + elif mode == "test": batch_out = [] for roidb in roidbs: - im, gt_boxes, gt_classes, is_crowd, im_info, im_id = roidb_reader( - roidb) - batch_out.append( - (im, gt_boxes, gt_classes, is_crowd, im_info, im_id)) + im, im_info, im_id = roidb_reader(roidb, mode) + batch_out.append((im, im_info, im_id)) if len(batch_out) == batch_size: yield batch_out batch_out = [] + if len(batch_out) != 0: + yield batch_out + + else: + for roidb in roidbs: + im, im_info, im_id = roidb_reader(roidb, mode) + batch_out = [(im, im_info, im_id)] + yield batch_out return reader @@ -163,3 +173,7 @@ def train(settings, def test(settings, batch_size, total_batch_size=None, padding_total=False): return coco(settings, 'test', batch_size, total_batch_size, shuffle=False) + + +def infer(settings): + return coco(settings, 'infer') diff --git a/fluid/faster_rcnn/roidbs.py b/fluid/faster_rcnn/roidbs.py index e7272733..d5d0d222 100644 --- a/fluid/faster_rcnn/roidbs.py +++ b/fluid/faster_rcnn/roidbs.py @@ -96,7 +96,8 @@ class JsonDataset(object): self._extend_with_flipped_entries(roidb) print('Loaded dataset: {:s}'.format(self.name)) print('{:d} roidb entries'.format(len(roidb))) - self._filter_for_training(roidb) + if self.is_train: + self._filter_for_training(roidb) return roidb def _prep_roidb_entry(self, entry): diff --git a/fluid/faster_rcnn/train.py b/fluid/faster_rcnn/train.py index 2c859318..ceec3efa 100644 --- a/fluid/faster_rcnn/train.py +++ b/fluid/faster_rcnn/train.py @@ -5,7 +5,7 @@ import argparse import functools import shutil import cPickle -from utility import add_arguments, print_arguments, SmoothedValue +from utility import parse_args, print_arguments, SmoothedValue import paddle import paddle.fluid as fluid @@ -14,40 +14,6 @@ import models.model_builder as model_builder import models.resnet as resnet from learning_rate import exponential_with_warmup_decay -parser = argparse.ArgumentParser(description=__doc__) -add_arg = functools.partial(add_arguments, argparser=parser) -# yapf: disable -# ENV -add_arg('parallel', bool, True, "Minibatch size.") -add_arg('use_gpu', bool, True, "Whether use GPU.") -add_arg('model_save_dir', str, 'output', "The path to save model.") -add_arg('pretrained_model', str, 'imagenet_resnet50_fusebn', "The init model path.") -add_arg('dataset', str, 'coco2017', "coco2014, coco2017, and pascalvoc.") -add_arg('data_dir', str, 'data/COCO17', "data directory") -add_arg('class_num', int, 81, "Class number.") -add_arg('use_pyreader', bool, True, "Use pyreader.") -add_arg('padding_minibatch',bool, False, - "If False, only resize image and not pad, image shape is different between" - " GPUs in one mini-batch. If True, image shape is the same in one mini-batch.") -# SOLVER -add_arg('learning_rate', float, 0.01, "Learning rate.") -add_arg('max_iter', int, 180000, "Iter number.") -add_arg('log_window', int, 1, "Log smooth window, set 1 for debug, set 20 for train.") -add_arg('snapshot_stride', int, 10000, "save model every snapshot stride.") -# RPN -add_arg('anchor_sizes', int, [32,64,128,256,512], "The size of anchors.") -add_arg('aspect_ratios', float, [0.5,1.0,2.0], "The ratio of anchors.") -add_arg('variance', float, [1.,1.,1.,1.], "The variance of anchors.") -add_arg('rpn_stride', float, 16., "Stride of the feature map that RPN is attached.") -# FAST RCNN -# TRAIN TEST -add_arg('batch_size', int, 8, "Minibatch size of all devices.") -add_arg('max_size', int, 1333, "The max resized image size.") -add_arg('scales', int, [800], "The resized image height.") -add_arg('batch_size_per_im',int, 512, "fast rcnn head batch size") -add_arg('mean_value', float, [102.9801, 115.9465, 122.7717], "pixel mean") -add_arg('debug', bool, False, "Debug mode") -#yapf: enable def train(cfg): learning_rate = cfg.learning_rate @@ -71,21 +37,22 @@ def train(cfg): use_random=True) model.build_model(image_shape) loss_cls, loss_bbox, rpn_cls_loss, rpn_reg_loss = model.loss() - loss_cls.persistable=True - loss_bbox.persistable=True - rpn_cls_loss.persistable=True - rpn_reg_loss.persistable=True + loss_cls.persistable = True + loss_bbox.persistable = True + rpn_cls_loss.persistable = True + rpn_reg_loss.persistable = True loss = loss_cls + loss_bbox + rpn_cls_loss + rpn_reg_loss boundaries = [120000, 160000] - values = [learning_rate, learning_rate*0.1, learning_rate*0.01] + values = [learning_rate, learning_rate * 0.1, learning_rate * 0.01] optimizer = fluid.optimizer.Momentum( - learning_rate=exponential_with_warmup_decay(learning_rate=learning_rate, + learning_rate=exponential_with_warmup_decay( + learning_rate=learning_rate, boundaries=boundaries, values=values, warmup_iter=500, - warmup_factor=1.0/3.0), + warmup_factor=1.0 / 3.0), regularization=fluid.regularizer.L2Decay(0.0001), momentum=0.9) optimizer.minimize(loss) @@ -97,10 +64,11 @@ def train(cfg): exe.run(fluid.default_startup_program()) if cfg.pretrained_model: + def if_exist(var): return os.path.exists(os.path.join(cfg.pretrained_model, var.name)) - fluid.io.load_vars(exe, cfg.pretrained_model, predicate=if_exist) + fluid.io.load_vars(exe, cfg.pretrained_model, predicate=if_exist) if cfg.parallel: train_exe = fluid.ParallelExecutor( @@ -109,17 +77,19 @@ def train(cfg): assert cfg.batch_size % devices_num == 0 batch_size_per_dev = cfg.batch_size / devices_num if cfg.use_pyreader: - train_reader = reader.train(cfg, batch_size=batch_size_per_dev, - total_batch_size=cfg.batch_size, - padding_total=cfg.padding_minibatch, - shuffle=True) + train_reader = reader.train( + cfg, + batch_size=batch_size_per_dev, + total_batch_size=cfg.batch_size, + padding_total=cfg.padding_minibatch, + shuffle=True) py_reader = model.py_reader py_reader.decorate_paddle_reader(train_reader) else: - train_reader = reader.train(cfg, batch_size=cfg.batch_size, shuffle=True) + train_reader = reader.train( + cfg, batch_size=cfg.batch_size, shuffle=True) feeder = fluid.DataFeeder(place=place, feed_list=model.feeds()) - def save_model(postfix): model_path = os.path.join(cfg.model_save_dir, postfix) if os.path.isdir(model_path): @@ -128,7 +98,7 @@ def train(cfg): fetch_list = [loss, rpn_cls_loss, rpn_reg_loss, loss_cls, loss_bbox] - def train_step_pyreader(): + def train_loop_pyreader(): py_reader.start() smoothed_loss = SmoothedValue(cfg.log_window) try: @@ -141,9 +111,12 @@ def train(cfg): losses = train_exe.run(fetch_list=[v.name for v in fetch_list]) every_pass_loss.append(np.mean(np.array(losses[0]))) smoothed_loss.add_value(np.mean(np.array(losses[0]))) - lr = np.array(fluid.global_scope().find_var('learning_rate').get_tensor()) + lr = np.array(fluid.global_scope().find_var('learning_rate') + .get_tensor()) print("Iter {:d}, lr {:.6f}, loss {:.6f}, time {:.5f}".format( - iter_id, lr[0], smoothed_loss.get_median_value(), start_time - prev_start_time)) + iter_id, lr[0], + smoothed_loss.get_median_value( + ), start_time - prev_start_time)) #print('cls_loss ', losses[1][0], ' reg_loss ', losses[2][0], ' loss_cls ', losses[3][0], ' loss_bbox ', losses[4][0]) if (iter_id + 1) % cfg.snapshot_stride == 0: save_model("model_iter{}".format(iter_id)) @@ -151,7 +124,7 @@ def train(cfg): py_reader.reset() return np.mean(every_pass_loss) - def train_step(): + def train_loop(): start_time = time.time() prev_start_time = start_time start = start_time @@ -165,9 +138,11 @@ def train(cfg): loss_v = np.mean(np.array(losses[0])) every_pass_loss.append(loss_v) smoothed_loss.add_value(loss_v) - lr = np.array(fluid.global_scope().find_var('learning_rate').get_tensor()) + lr = np.array(fluid.global_scope().find_var('learning_rate') + .get_tensor()) print("Iter {:d}, lr {:.6f}, loss {:.6f}, time {:.5f}".format( - iter_id, lr[0], smoothed_loss.get_median_value(), start_time - prev_start_time)) + iter_id, lr[0], + smoothed_loss.get_median_value(), start_time - prev_start_time)) #print('cls_loss ', losses[1][0], ' reg_loss ', losses[2][0], ' loss_cls ', losses[3][0], ' loss_bbox ', losses[4][0]) if (iter_id + 1) % cfg.snapshot_stride == 0: save_model("model_iter{}".format(iter_id)) @@ -176,13 +151,14 @@ def train(cfg): return np.mean(every_pass_loss) if cfg.use_pyreader: - train_step_pyreader() + train_loop_pyreader() else: - train_step() + train_loop() save_model('model_final') + if __name__ == '__main__': - args = parser.parse_args() + args = parse_args() print_arguments(args) data_args = reader.Settings(args) diff --git a/fluid/faster_rcnn/utility.py b/fluid/faster_rcnn/utility.py index ca4c00da..8f8f6410 100644 --- a/fluid/faster_rcnn/utility.py +++ b/fluid/faster_rcnn/utility.py @@ -22,6 +22,8 @@ import numpy as np import six from collections import deque from paddle.fluid import core +import argparse +import functools def print_arguments(args): @@ -78,3 +80,51 @@ class SmoothedValue(object): def get_median_value(self): return np.median(self.deque) + + +def parse_args(): + """ + return all args + """ + parser = argparse.ArgumentParser(description=__doc__) + add_arg = functools.partial(add_arguments, argparser=parser) + # yapf: disable + # ENV + add_arg('parallel', bool, True, "Whether use parallel.") + add_arg('use_gpu', bool, True, "Whether use GPU.") + add_arg('model_save_dir', str, 'output', "The path to save model.") + add_arg('pretrained_model', str, 'imagenet_resnet50_fusebn', "The init model path.") + add_arg('dataset', str, 'coco2017', "coco2014, coco2017.") + add_arg('data_dir', str, 'data/COCO17', "The data root path.") + add_arg('class_num', int, 81, "Class number.") + add_arg('use_pyreader', bool, True, "Use pyreader.") + add_arg('padding_minibatch',bool, False, + "If False, only resize image and not pad, image shape is different between" + " GPUs in one mini-batch. If True, image shape is the same in one mini-batch.") + #SOLVER + add_arg('learning_rate', float, 0.01, "Learning rate.") + add_arg('max_iter', int, 180000, "Iter number.") + add_arg('log_window', int, 1, "Log smooth window, set 1 for debug, set 20 for train.") + add_arg('snapshot_stride', int, 10000, "save model every snapshot stride.") + # RPN + add_arg('anchor_sizes', int, [32,64,128,256,512], "The size of anchors.") + add_arg('aspect_ratios', float, [0.5,1.0,2.0], "The ratio of anchors.") + add_arg('variance', float, [1.,1.,1.,1.], "The variance of anchors.") + add_arg('rpn_stride', float, 16., "Stride of the feature map that RPN is attached.") + # FAST RCNN + # TRAIN TEST INFER + add_arg('batch_size', int, 1, "Minibatch size.") + add_arg('max_size', int, 1333, "The resized image height.") + add_arg('scales', int, [800], "The resized image height.") + add_arg('batch_size_per_im',int, 512, "fast rcnn head batch size") + add_arg('mean_value', float, [102.9801, 115.9465, 122.7717], "pixel mean") + add_arg('nms_threshold', float, 0.5, "NMS threshold.") + add_arg('score_threshold', float, 0.05, "score threshold for NMS.") + add_arg('debug', bool, False, "Debug mode") + # SINGLE EVAL AND DRAW + add_arg('draw_threshold', float, 0.8, "Confidence threshold to draw bbox.") + add_arg('image_path', str, 'data/COCO17/val2017', "The image path used to inference and visualize.") + add_arg('image_name', str, '', "The single image used to inference and visualize.") + # yapf: enable + args = parser.parse_args() + return args -- GitLab