From f8cddad5e7f8e151e8b02131c570eac93ea9db1f Mon Sep 17 00:00:00 2001
From: tangwei12 <tangwei12@baidu.com>
Date: Thu, 22 Nov 2018 17:06:06 +0800
Subject: [PATCH] rewrite pyreader

---
 fluid/PaddleRec/word2vec/network_conf.py | 26 +++++++++++++-----------
 fluid/PaddleRec/word2vec/train.py        |  5 ++++-
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/fluid/PaddleRec/word2vec/network_conf.py b/fluid/PaddleRec/word2vec/network_conf.py
index f91ecab9..f9eb4ca5 100644
--- a/fluid/PaddleRec/word2vec/network_conf.py
+++ b/fluid/PaddleRec/word2vec/network_conf.py
@@ -96,31 +96,33 @@ def skip_gram_word2vec(dict_size,
     data_lod_levels.append(1)
     data_types.append('int64')
 
-    py_reader = fluid.layers.py_reader(capacity=64,
-                                       shapes=data_shapes,
-                                       lod_levels=data_lod_levels,
-                                       dtypes=data_types,
-                                       name='py_reader',
-                                       use_double_buffer=True)
+    datas = []
 
-    word_and_label = fluid.layers.read_file(py_reader)
+    input_word = fluid.layers.data(name="input_word", shape=[1], dtype='int64')
+    predict_word = fluid.layers.data(name='predict_word', shape=[1], dtype='int64')
+
+    datas.append(input_word, predict_word)
 
     cost = None
 
     emb = fluid.layers.embedding(
-        input=word_and_label[0],
+        input=input_word,
         is_sparse=is_sparse,
         size=[dict_size, embedding_size],
         param_attr=fluid.ParamAttr(initializer=fluid.initializer.Normal(
             scale=1 / math.sqrt(dict_size))))
 
     if with_nce:
-        cost = nce_layer(emb, word_and_label[1], embedding_size, dict_size, 5,
-                         "uniform", word_frequencys, None)
+        cost = nce_layer(emb, predict_word, embedding_size, dict_size, 5, "uniform",
+                         word_frequencys, None)
     if with_hsigmoid:
-        cost = hsigmoid_layer(emb, word_and_label[1], dict_size, max_code_length,
-                              None)
+        cost = hsigmoid_layer(emb, predict_word, dict_size, max_code_length, datas)
 
     avg_cost = fluid.layers.reduce_mean(cost)
 
+    py_reader = fluid.layers.create_py_reader_by_data(capacity=64,
+                                                      feed_list=datas,
+                                                      name='py_reader',
+                                                      use_double_buffer=True)
+
     return avg_cost, py_reader
diff --git a/fluid/PaddleRec/word2vec/train.py b/fluid/PaddleRec/word2vec/train.py
index 230a3421..8b9417dc 100644
--- a/fluid/PaddleRec/word2vec/train.py
+++ b/fluid/PaddleRec/word2vec/train.py
@@ -107,7 +107,10 @@ def train_loop(args, train_program, reader, py_reader, loss, trainer_id):
     start = time.clock()
 
     exec_strategy = fluid.ExecutionStrategy()
-    exec_strategy.num_threads = int(os.getenv("NUM_THREADS"))
+
+    if os.getenv("NUM_THREADS", ""):
+        exec_strategy.num_threads = int(os.getenv("NUM_THREADS"))
+
     build_strategy = fluid.BuildStrategy()
     build_strategy.reduce_strategy = fluid.BuildStrategy.ReduceStrategy.Reduce
 
-- 
GitLab