Refine code

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

@@ -12,9 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 import paddle.v2 as paddle
 import paddle.v2.fluid as fluid
+import unittest
+
+
+def main_impl(use_cuda):
+    if use_cuda and not fluid.core.is_compiled_with_cuda():
+        return
 
     PASS_NUM = 100
     EMBED_SIZE = 32
@@ -59,7 +64,9 @@ embed_forth = fluid.layers.embedding(
 
     concat_embed = fluid.layers.concat(
         input=[embed_first, embed_second, embed_third, embed_forth], axis=1)
-hidden1 = fluid.layers.fc(input=concat_embed, size=HIDDEN_SIZE, act='sigmoid')
+    hidden1 = fluid.layers.fc(input=concat_embed,
+                              size=HIDDEN_SIZE,
+                              act='sigmoid')
     predict_word = fluid.layers.fc(input=hidden1, size=dict_size, act='softmax')
     cost = fluid.layers.cross_entropy(input=predict_word, label=next_word)
     avg_cost = fluid.layers.mean(x=cost)
@@ -69,7 +76,7 @@ sgd_optimizer.minimize(avg_cost)
     train_reader = paddle.batch(
         paddle.dataset.imikolov.train(word_dict, N), BATCH_SIZE)
 
-place = fluid.CPUPlace()
+    place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
     exe = fluid.Executor(place)
     feeder = fluid.DataFeeder(
         feed_list=[first_word, second_word, third_word, forth_word, next_word],
@@ -83,5 +90,26 @@ for pass_id in range(PASS_NUM):
                                   feed=feeder.feed(data),
                                   fetch_list=[avg_cost])
             if avg_cost_np[0] < 5.0:
-            exit(0)  # if avg cost less than 10.0, we think our code is good.
-exit(1)
+                return
+    raise AssertionError("Cost is too large {0:2.2}".format(avg_cost_np[0]))
+
+
+def main(*args, **kwargs):
+    prog = fluid.Program()
+    startup_prog = fluid.Program()
+    scope = fluid.core.Scope()
+    with fluid.scope_guard(scope):
+        with fluid.program_guard(prog, startup_prog):
+            main_impl(*args, **kwargs)
+
+
+class W2VTest(unittest.TestCase):
+    def test_cpu_normal(self):
+        main(use_cuda=False)
+
+    def test_gpu_normal(self):
+        main(use_cuda=True)
+
+
+if __name__ == '__main__':
+    unittest.main()