src/caffe/test/data/mnist.py

   1 """This script generates the mnist train and test leveldbs used in the
   2 test.
   3 """
   4 from caffe.pyutil import convert
   5 from decaf.layers import core_layers
   6 import numpy as np
   7 import leveldb
   8
   9 # the folder that has the MNIST data
  10 MNIST_ROOT = 'mnist'
  11
  12 mnist = core_layers.MNISTDataLayer(
  13     rootfolder=MNIST_ROOT, name='mnist', is_training = True)
  14 db = leveldb.LevelDB('mnist-train-leveldb')
  15
  16 for i in range(60000):
  17   datum = convert.array_to_datum((mnist._data[i] * 255).reshape(1,28,28).astype(np.uint8))
  18   datum.label = mnist._label[i]
  19   db.Put('%d' % (i), datum.SerializeToString())
  20 del db