Merge pull request #881 from kuke/refactor_mobilenet

Enable the parallel training of mobilenet

fluid/image_classification/mobilenet.py

@@ -153,72 +153,3 @@ def mobile_net(img, class_dim, scale=1.0):
                           act='softmax',
                           param_attr=parameter_attr)
     return tmp
-
-
-def train(learning_rate, batch_size, num_passes, model_save_dir='model'):
-    class_dim = 102
-    image_shape = [3, 224, 224]
-
-    image = fluid.layers.data(name='image', shape=image_shape, dtype='float32')
-    label = fluid.layers.data(name='label', shape=[1], dtype='int64')
-
-    out = mobile_net(image, class_dim=class_dim)
-
-    cost = fluid.layers.cross_entropy(input=out, label=label)
-    avg_cost = fluid.layers.mean(x=cost)
-
-    optimizer = fluid.optimizer.Momentum(
-        learning_rate=learning_rate,
-        momentum=0.9,
-        regularization=fluid.regularizer.L2Decay(5 * 1e-5))
-    opts = optimizer.minimize(avg_cost)
-
-    b_size_var = fluid.layers.create_tensor(dtype='int64')
-    b_acc_var = fluid.layers.accuracy(input=out, label=label, total=b_size_var)
-
-    inference_program = fluid.default_main_program().clone()
-    with fluid.program_guard(inference_program):
-        inference_program = fluid.io.get_inference_program(
-            target_vars=[b_acc_var, b_size_var])
-
-    place = fluid.CPUPlace()
-    exe = fluid.Executor(place)
-    exe.run(fluid.default_startup_program())
-
-    train_reader = paddle.batch(
-        paddle.dataset.flowers.train(), batch_size=batch_size)
-    test_reader = paddle.batch(
-        paddle.dataset.flowers.test(), batch_size=batch_size)
-    feeder = fluid.DataFeeder(place=place, feed_list=[image, label])
-
-    train_pass_acc_evaluator = fluid.average.WeightedAverage()
-    test_pass_acc_evaluator = fluid.average.WeightedAverage()
-    for pass_id in range(num_passes):
-        train_pass_acc_evaluator.reset()
-        for batch_id, data in enumerate(train_reader()):
-            loss, acc, size = exe.run(
-                fluid.default_main_program(),
-                feed=feeder.feed(data),
-                fetch_list=[avg_cost, b_acc_var, b_size_var])
-            train_pass_acc_evaluator.add(value=acc, weight=size)
-            print("Pass {0}, batch {1}, loss {2}, acc {3}".format(
-                pass_id, batch_id, loss[0], acc[0]))
-
-        test_pass_acc_evaluator.reset()
-        for data in test_reader():
-            loss, acc, size = exe.run(
-                inference_program,
-                feed=feeder.feed(data),
-                fetch_list=[avg_cost, b_acc_var, b_size_var])
-            test_pass_acc_evaluator.add(value=acc, weight=size)
-        print("End pass {0}, train_acc {1}, test_acc {2}".format(
-            pass_id,
-            train_pass_acc_evaluator.eval(), test_pass_acc_evaluator.eval()))
-        if pass_id % 10 == 0:
-            model_path = os.path.join(model_save_dir, str(pass_id))
-            print 'save models to %s' % (model_path)
-            fluid.io.save_inference_model(model_path, ['image'], [out], exe)
-
-
-if __name__ == '__main__':
-    train(learning_rate=0.005, batch_size=40, num_passes=300)

fluid/image_classification/train.py

@@ -5,6 +5,7 @@ import sys
 import paddle.v2 as paddle
 import paddle.fluid as fluid
 from se_resnext import SE_ResNeXt
+from mobilenet import mobile_net
 import reader
 
 import argparse
@@ -18,6 +19,11 @@ add_arg('batch_size',   int,  256, "Minibatch size.")
 add_arg('num_layers',   int,  50,  "How many layers for SE-ResNeXt model.")
 add_arg('with_mem_opt', bool, True, "Whether to use memory optimization or not.")
 add_arg('parallel_exe', bool, True, "Whether to use ParallelExecutor to train or not.")
+add_arg('model',
+        type=str,
+        choices=['mobilenet', 'se_resnext'],
+        default='se_resnext',
+        help="Which model to use.")
 # yapf: enable
 
 
@@ -44,7 +50,12 @@ def train_parallel_do(args,
         with pd.do():
             image_ = pd.read_input(image)
             label_ = pd.read_input(label)
-            out = SE_ResNeXt(input=image_, class_dim=class_dim, layers=layers)
+            if args.model is 'se_resnext':
+                out = SE_ResNeXt(
+                    input=image_, class_dim=class_dim, layers=layers)
+            else:
+                out = mobile_net(img=image_, class_dim=class_dim)
+
             cost = fluid.layers.cross_entropy(input=out, label=label_)
             avg_cost = fluid.layers.mean(x=cost)
             acc_top1 = fluid.layers.accuracy(input=out, label=label_, k=1)
@@ -58,7 +69,11 @@ def train_parallel_do(args,
         acc_top1 = fluid.layers.mean(x=acc_top1)
         acc_top5 = fluid.layers.mean(x=acc_top5)
     else:
+        if args.model is 'se_resnext':
             out = SE_ResNeXt(input=image, class_dim=class_dim, layers=layers)
+        else:
+            out = mobile_net(img=image, class_dim=class_dim)
+
         cost = fluid.layers.cross_entropy(input=out, label=label)
         avg_cost = fluid.layers.mean(x=cost)
         acc_top1 = fluid.layers.accuracy(input=out, label=label, k=1)
@@ -150,7 +165,8 @@ def train_parallel_do(args,
               test_acc5))
         sys.stdout.flush()
 
-        model_path = os.path.join(model_save_dir, str(pass_id))
+        model_path = os.path.join(model_save_dir + '/' + args.model,
+                                  str(pass_id))
         if not os.path.isdir(model_path):
             os.makedirs(model_path)
         fluid.io.save_persistables(exe, model_path)
@@ -171,7 +187,11 @@ def train_parallel_exe(args,
 
     image = fluid.layers.data(name='image', shape=image_shape, dtype='float32')
     label = fluid.layers.data(name='label', shape=[1], dtype='int64')
+    if args.model is 'se_resnext':
         out = SE_ResNeXt(input=image, class_dim=class_dim, layers=layers)
+    else:
+        out = mobile_net(img=image, class_dim=class_dim)
+
     cost = fluid.layers.cross_entropy(input=out, label=label)
     acc_top1 = fluid.layers.accuracy(input=out, label=label, k=1)
     acc_top5 = fluid.layers.accuracy(input=out, label=label, k=5)
@@ -270,7 +290,8 @@ def train_parallel_exe(args,
               test_acc5))
         sys.stdout.flush()
 
-        model_path = os.path.join(model_save_dir, str(pass_id))
+        model_path = os.path.join(model_save_dir + '/' + args.model,
+                                  str(pass_id))
         if not os.path.isdir(model_path):
             os.makedirs(model_path)
         fluid.io.save_persistables(exe, model_path)