# coding: utf-8

from pylab import *

caffe_root = '/Users/wang/soft/caffe/'  # this file should be run from {caffe_root}/examples (otherwise change this line)

import sys

sys.path.insert(0, caffe_root + 'python')
import caffe

from caffe import layers as L, params as P


def lenet(lmdb, batch_size):
    n = caffe.NetSpec()
    n.data, n.label = L.Data(
        batch_size=batch_size,
        backend=P.Data.LMDB,
        source=lmdb,
        transform_param=dict(scale=1. / 256),
        ntop=2
    )
    n.conv1 = L.Convolution(n.data, kernel_size=5, num_output=20, weight_filler=dict(type='xavier'))
    n.pool1 = L.Pooling(n.conv1, kernel_size=2, stride=2, pool=P.Pooling.MAX)
    n.conv2 = L.Convolution(n.pool1, kernel_size=5, num_output=50, weight_filler=dict(type='xavier'))
    n.pool2 = L.Pooling(n.conv2, kernel_size=2, stride=2, pool=P.Pooling.MAX)
    n.fc1 = L.InnerProduct(n.pool2, num_output=500, weight_filler=dict(type='xavier'))
    n.relu1 = L.ReLU(n.fc1, in_place=True)
    n.score = L.InnerProduct(n.relu1, num_output=10, weight_filler=dict(type='xavier'))
    n.loss = L.SoftmaxWithLoss(n.score, n.label)

    return n.to_proto()

my_lenet = lenet(caffe_root + 'examples/mnist/mnist_train_lmdb', 64)

print my_lenet