add infer (#1274)

* refine test reader return

* add infer and parse_args

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)

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)

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,

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

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

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)

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