Merge pull request #774 from qingqing01/ssd

Refine train.py for MobileNet-SSD.

fluid/object_detection/image_util.py

@@ -216,7 +216,7 @@ def distort_image(img, settings):
 
 def expand_image(img, bbox_labels, img_width, img_height, settings):
     prob = random.uniform(0, 1)
-    if prob < settings._hue_prob:
+    if prob < settings._expand_prob:
         expand_ratio = random.uniform(1, settings._expand_max_ratio)
         if expand_ratio - 1 >= 0.01:
             height = int(img_height * expand_ratio)

fluid/object_detection/train.py

@@ -12,8 +12,9 @@ import functools
 parser = argparse.ArgumentParser(description=__doc__)
 add_arg = functools.partial(add_arguments, argparser=parser)
 # yapf: disable
-add_arg('parallel',    bool,   True,     "Whether use parallel training.")
-add_arg('use_gpu',     bool,   True,     "Whether use GPU.")
+add_arg('batch_size',   int,    32,       "Minibatch size.")
+add_arg('parallel',     bool,   True,     "Whether use parallel training.")
+add_arg('use_gpu',      bool,   True,     "Whether use GPU.")
 # yapf: disable
 
 
@@ -47,26 +48,24 @@ def train(args,
             locs, confs, box, box_var = mobile_net(image_, image_shape)
             loss = fluid.layers.ssd_loss(locs, confs, gt_box_, gt_label_,
                                          box, box_var)
+            nmsed_out = fluid.layers.detection_output(
+                locs, confs, box, box_var, nms_threshold=0.45)
+            loss = fluid.layers.reduce_sum(loss)
             pd.write_output(loss)
-            pd.write_output(locs)
-            pd.write_output(confs)
-            pd.write_output(box)
-            pd.write_output(box_var)
+            pd.write_output(nmsed_out)
 
-        loss, locs, confs, box, box_var = pd()
-        loss = fluid.layers.reduce_sum(loss)
+        loss, nmsed_out = pd()
+        loss = fluid.layers.mean(loss)
     else:
         locs, confs, box, box_var = mobile_net(image, image_shape)
         nmsed_out = fluid.layers.detection_output(
-            locs, mbox_confs, box, box_var, nms_threshold=0.45)
-        loss = fluid.layers.ssd_loss(locs, mbox_confs, gt_box, gt_label,
+            locs, confs, box, box_var, nms_threshold=0.45)
+        loss = fluid.layers.ssd_loss(locs, confs, gt_box, gt_label,
                                      box, box_var)
         loss = fluid.layers.reduce_sum(loss)
 
     test_program = fluid.default_main_program().clone(for_test=True)
     with fluid.program_guard(test_program):
-        nmsed_out = fluid.layers.detection_output(
-            locs, confs, box, box_var, nms_threshold=0.45)
         map_eval = fluid.evaluator.DetectionMAP(
             nmsed_out,
             gt_label,
@@ -98,7 +97,6 @@ def train(args,
     feeder = fluid.DataFeeder(
         place=place, feed_list=[image, gt_box, gt_label, difficult])
 
-    #print 'test_program ', test_program
     def test(pass_id):
         _, accum_map = map_eval.get_map_var()
         map_eval.reset(exe)
@@ -109,7 +107,6 @@ def train(args,
                                fetch_list=[accum_map])
         print("Test {0}, map {1}".format(pass_id, test_map[0]))
 
-    #print 'main_program ', fluid.default_main_program()
     for pass_id in range(num_passes):
         for batch_id, data in enumerate(train_reader()):
             loss_v = exe.run(fluid.default_main_program(),
@@ -143,5 +140,5 @@ if __name__ == '__main__':
           val_file_list='./data/test.txt',
           data_args=data_args,
           learning_rate=0.001,
-          batch_size=32,
+          batch_size=args.batch_size,
           num_passes=300)