From 283bdc5062be0ba14b0ae3ca6cc211ddaf25fd1c Mon Sep 17 00:00:00 2001
From: gongweibao <weibao.gong@gmail.com>
Date: Mon, 12 Jun 2017 10:29:35 +0800
Subject: [PATCH] fix by helin's comments

---
 paddle/parameter/tests/test_argument.cpp      |  2 +-
 python/paddle/v2/dataset/common.py            | 58 +++++++++++--------
 python/paddle/v2/dataset/tests/common_test.py | 26 +++++++--
 3 files changed, 56 insertions(+), 30 deletions(-)

diff --git a/paddle/parameter/tests/test_argument.cpp b/paddle/parameter/tests/test_argument.cpp
index 81fe4ee397..98ab013548 100644
--- a/paddle/parameter/tests/test_argument.cpp
+++ b/paddle/parameter/tests/test_argument.cpp
@@ -42,7 +42,7 @@ TEST(Argument, poolSequenceWithStride) {
     CHECK_EQ(outStart[3], 4);
     CHECK_EQ(outStart[4], 7);
 
-    CHECK_EQ(stridePositions->getSize(), 8);
+    CHECK_EQ(stridePositions->getSize(), 8UL);
     auto result = reversed ? strideResultReversed : strideResult;
     for (int i = 0; i < 8; i++) {
       CHECK_EQ(stridePositions->getData()[i], result[i]);
diff --git a/python/paddle/v2/dataset/common.py b/python/paddle/v2/dataset/common.py
index 89675080e2..8023fa3cf8 100644
--- a/python/paddle/v2/dataset/common.py
+++ b/python/paddle/v2/dataset/common.py
@@ -151,9 +151,14 @@ def cluster_files_reader(files_pattern,
     return reader
 
 
-def convert(output_path, eader, num_shards, name_prefix):
+def convert(output_path,
+            reader,
+            num_shards,
+            name_prefix,
+            max_lines_to_shuffle=10000):
     import recordio
     import cPickle as pickle
+    import random
     """
     Convert data from reader to recordio format files.
 
@@ -161,35 +166,40 @@ def convert(output_path, eader, num_shards, name_prefix):
     :param reader: a data reader, from which the convert program will read data instances.
     :param num_shards: the number of shards that the dataset will be partitioned into.
     :param name_prefix: the name prefix of generated files.
+    :param max_lines_to_shuffle: the max lines numbers to shuffle before writing.
     """
 
-    def open_needs(idx):
-        n = "%s/%s-%05d" % (output_path, name_prefix, idx)
-        w = recordio.writer(n)
-        f = open(n, "w")
-        idx += 1
+    assert num_shards >= 1
+    assert max_lines_to_shuffle >= 1
 
-        return w, f, idx
+    def open_writers():
+        w = []
+        for i in range(0, num_shards):
+            n = "%s/%s-%05d-of-%05d" % (output_path, name_prefix, i,
+                                        num_shards - 1)
+            w.append(recordio.writer(n))
 
-    def close_needs(w, f):
-        if w is not None:
-            w.close()
+        return w
 
-        if f is not None:
-            f.close()
+    def close_writers(w):
+        for i in range(0, num_shards):
+            w[i].close()
 
-    idx = 0
-    w = None
-    f = None
+    def write_data(w, lines):
+        random.shuffle(lines)
+        for i, d in enumerate(lines):
+            d = pickle.dumps(d, pickle.HIGHEST_PROTOCOL)
+            w[i % num_shards].write(d)
 
-    for i, d in enumerate(reader()):
-        if w is None:
-            w, f, idx = open_needs(idx)
-
-        w.write(pickle.dumps(d, pickle.HIGHEST_PROTOCOL))
+    w = open_writers()
+    lines = []
 
-        if i % num_shards == 0 and i >= num_shards:
-            close_needs(w, f)
-            w, f, idx = open_needs(idx)
+    for i, d in enumerate(reader()):
+        lines.append(d)
+        if i % max_lines_to_shuffle == 0 and i >= max_lines_to_shuffle:
+            write_data(w, lines)
+            lines = []
+            continue
 
-    close_needs(w, f)
+    write_data(w, lines)
+    close_writers(w)
diff --git a/python/paddle/v2/dataset/tests/common_test.py b/python/paddle/v2/dataset/tests/common_test.py
index 3120026e1e..cfa194eba3 100644
--- a/python/paddle/v2/dataset/tests/common_test.py
+++ b/python/paddle/v2/dataset/tests/common_test.py
@@ -58,20 +58,36 @@ class TestCommon(unittest.TestCase):
             self.assertEqual(e, str("0"))
 
     def test_convert(self):
+        record_num = 10
+        num_shards = 4
+
         def test_reader():
             def reader():
-                for x in xrange(10):
+                for x in xrange(record_num):
                     yield x
 
             return reader
 
         path = tempfile.mkdtemp()
-
         paddle.v2.dataset.common.convert(path,
-                                         test_reader(), 4, 'random_images')
+                                         test_reader(), num_shards,
+                                         'random_images')
 
-        files = glob.glob(temp_path + '/random_images-*')
-        self.assertEqual(len(files), 3)
+        files = glob.glob(path + '/random_images-*')
+        self.assertEqual(len(files), num_shards)
+
+        recs = []
+        for i in range(0, num_shards):
+            n = "%s/random_images-%05d-of-%05d" % (path, i, num_shards - 1)
+            r = recordio.reader(n)
+            while True:
+                d = r.read()
+                if d is None:
+                    break
+                recs.append(d)
+
+        recs.sort()
+        self.assertEqual(total, record_num)
 
 
 if __name__ == '__main__':
-- 
GitLab