MNIST dataset reader implementation

python/paddle/v2/data_set/mnist.py

+import sklearn.datasets.mldata
+import sklearn.model_selection
+import numpy
+
+__all__ = ['MNISTReader', 'train_reader_creator', 'test_reader_creator']
+
+DATA_HOME = None
+
+
+def __mnist_reader__(data, target):
+    n_samples = data.shape[0]
+    for i in xrange(n_samples):
+        yield data[i].astype(numpy.float32), int(target[i])
+
+
+class MNISTReader(object):
+    """
+    mnist dataset reader. The `train_reader` and `test_reader` method returns
+    a iterator of each sample. Each sample is combined by 784-dim float and a
+    one-dim label
+    """
+
+    def __init__(self, random_state):
+        data = sklearn.datasets.mldata.fetch_mldata(
+            "MNIST original", data_home=DATA_HOME)
+        n_train = 60000
+        self.X_train, self.X_test, self.y_train, self.y_test = sklearn.model_selection.train_test_split(
+            data.data / 255.0,
+            data.target.astype("int"),
+            train_size=n_train,
+            random_state=random_state)
+
+    def train_reader(self):
+        return __mnist_reader__(self.X_train, self.y_train)
+
+    def test_reader(self):
+        return __mnist_reader__(self.X_test, self.y_test)
+
+
+__default_instance__ = MNISTReader(0)
+
+
+def train_reader_creator():
+    """
+    Default train set reader creator.
+    """
+    return __default_instance__.train_reader
+
+
+def test_reader_creator():
+    """
+    Default test set reader creator.
+    """
+    return __default_instance__.test_reader
+
+
+def unittest():
+    assert len(list(train_reader_creator()())) == 60000
+
+
+if __name__ == '__main__':
+    unittest()