wrtie_mnist_data_to_lmdb_v2.py

from mnist import MNIST
import numpy as np
import lmdb
import caffe
mndata = MNIST('./data')

count = 0
env = lmdb.open('mnist_lmdb', map_size=1000*1000*1000)

txn = env.begin(write=True)
for image, label in mndata.load_training():
    datum = caffe.proto.caffe_pb2.Datum()
    datum.channels = 1
    datum.height = 28
    datum.width = 28
    datum.data = np.array(image, dtype=np.uint8).tobytes()
    datum.label = label
    str_id = '{:08}'.format(count)
    txn.put(str_id, datum.SerializeToString())

    count = count + 1

    if count % 1000 == 0:
        txn.commit()
        txn = env.begin(write=True)

if count % 1000 != 0:
    txn.commit()
env.close()