Make image_classification as a normal python unittest

python/paddle/v2/fluid/tests/book/CMakeLists.txt

 file(GLOB TEST_OPS RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "test_*.py")
 string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
 
-list(REMOVE_ITEM TEST_OPS test_image_classification_train test_recognize_digits)
-py_test(test_image_classification_train_resnet SRCS test_image_classification_train.py ARGS resnet)
-py_test(test_image_classification_train_vgg SRCS test_image_classification_train.py ARGS vgg)
+list(REMOVE_ITEM TEST_OPS test_recognize_digits)
 py_test(test_recognize_digits_mlp_cpu
   SRCS test_recognize_digits.py
   ARGS mlp)

python/paddle/v2/fluid/tests/book/test_image_classification_train.py

@@ -14,10 +14,10 @@
 
 from __future__ import print_function
 
-import sys
-
 import paddle.v2 as paddle
 import paddle.v2.fluid as fluid
+import unittest
+import contextlib
 
 
 def resnet_cifar10(input, depth=32):
@@ -89,48 +89,49 @@ def vgg16_bn_drop(input):
     return fc2
 
 
-classdim = 10
-data_shape = [3, 32, 32]
+def main(net_type, use_cuda):
+    if use_cuda and not fluid.core.is_compiled_with_cuda():
+        return
 
-images = fluid.layers.data(name='pixel', shape=data_shape, dtype='float32')
-label = fluid.layers.data(name='label', shape=[1], dtype='int64')
+    classdim = 10
+    data_shape = [3, 32, 32]
 
-net_type = "vgg"
-if len(sys.argv) >= 2:
-    net_type = sys.argv[1]
+    images = fluid.layers.data(name='pixel', shape=data_shape, dtype='float32')
+    label = fluid.layers.data(name='label', shape=[1], dtype='int64')
 
-if net_type == "vgg":
+    if net_type == "vgg":
         print("train vgg net")
         net = vgg16_bn_drop(images)
-elif net_type == "resnet":
+    elif net_type == "resnet":
         print("train resnet")
         net = resnet_cifar10(images, 32)
-else:
+    else:
         raise ValueError("%s network is not supported" % net_type)
 
-predict = fluid.layers.fc(input=net, size=classdim, act='softmax')
-cost = fluid.layers.cross_entropy(input=predict, label=label)
-avg_cost = fluid.layers.mean(x=cost)
+    predict = fluid.layers.fc(input=net, size=classdim, act='softmax')
+    cost = fluid.layers.cross_entropy(input=predict, label=label)
+    avg_cost = fluid.layers.mean(x=cost)
 
-optimizer = fluid.optimizer.Adam(learning_rate=0.001)
-opts = optimizer.minimize(avg_cost)
+    optimizer = fluid.optimizer.Adam(learning_rate=0.001)
+    optimizer.minimize(avg_cost)
 
-accuracy = fluid.evaluator.Accuracy(input=predict, label=label)
+    accuracy = fluid.evaluator.Accuracy(input=predict, label=label)
 
-BATCH_SIZE = 128
-PASS_NUM = 1
+    BATCH_SIZE = 128
+    PASS_NUM = 1
 
-train_reader = paddle.batch(
+    train_reader = paddle.batch(
         paddle.reader.shuffle(
             paddle.dataset.cifar.train10(), buf_size=128 * 10),
         batch_size=BATCH_SIZE)
 
-place = fluid.CPUPlace()
-exe = fluid.Executor(place)
-feeder = fluid.DataFeeder(place=place, feed_list=[images, label])
-exe.run(fluid.default_startup_program())
+    place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
+    exe = fluid.Executor(place)
+    feeder = fluid.DataFeeder(place=place, feed_list=[images, label])
+    exe.run(fluid.default_startup_program())
 
-for pass_id in range(PASS_NUM):
+    loss = 0.0
+    for pass_id in range(PASS_NUM):
         accuracy.reset(exe)
         for data in train_reader():
             loss, acc = exe.run(fluid.default_main_program(),
@@ -139,6 +140,38 @@ for pass_id in range(PASS_NUM):
             pass_acc = accuracy.eval(exe)
             print("loss:" + str(loss) + " acc:" + str(acc) + " pass_acc:" + str(
                 pass_acc))
-        # this model is slow, so if we can train two mini batch, we think it works properly.
-        exit(0)
-exit(1)
+            return
+
+    raise AssertionError(
+        "Image classification loss is too large, {0:2.2}".format(loss))
+
+
+class TestImageClassification(unittest.TestCase):
+    def test_vgg_cuda(self):
+        with self.scope_prog_guard():
+            main('vgg', use_cuda=True)
+
+    def test_resnet_cuda(self):
+        with self.scope_prog_guard():
+            main('resnet', use_cuda=True)
+
+    def test_vgg_cpu(self):
+        with self.scope_prog_guard():
+            main('vgg', use_cuda=False)
+
+    def test_resnet_cpu(self):
+        with self.scope_prog_guard():
+            main('resnet', use_cuda=False)
+
+    @contextlib.contextmanager
+    def scope_prog_guard(self):
+        prog = fluid.Program()
+        startup_prog = fluid.Program()
+        scope = fluid.core.Scope()
+        with fluid.scope_guard(scope):
+            with fluid.program_guard(prog, startup_prog):
+                yield
+
+
+if __name__ == '__main__':
+    unittest.main()