from __future__ import print_function

import mxnet as mx

from mxnet import nd, autograd

from mxnet import gluon

import numpy as np

mx.random.seed(1)

ctx = mx.cpu()

mnist = mx.test_utils.get_mnist()

batch_size = 64

num_inputs = 784

num_outputs = 10

train_data = mx.io.NDArrayIter(mnist["train_data"], mnist["train_label"], batch_size, shuffle=True)

test_data = mx.io.NDArrayIter(mnist["test_data"], mnist["test_label"], batch_size, shuffle=True)

num_fc = 512

net = gluon.nn.Sequential()

with net.name_scope():

net.add(gluon.nn.Conv2D(channels=20, kernel_size=5, activation='relu'))

net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))

net.add(gluon.nn.Conv2D(channels=50, kernel_size=5, activation='relu'))

net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))

net.add(gluon.nn.Flatten())

net.add(gluon.nn.Dense(num_fc, activation="relu"))

net.add(gluon.nn.Dense(num_outputs))

net.collect_params().initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)

softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss()

trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': .1})

def evaluate_accuracy(data_iterator, net):

return acc.get()[1]

epochs = 10

smoothing_constant = .01

for e in range(epochs):

train_data.reset()

for i, batch in enumerate(train_data):

data = batch.data[0].as_in_context(ctx)

label = batch.label[0].as_in_context(ctx)

with autograd.record():

output = net(data)

loss = softmax_cross_entropy(output, label)

loss.backward()

trainer.step(data.shape[0])

##########################

# Keep a moving average of the losses

##########################

curr_loss = nd.mean(loss).asscalar()

moving_loss = (curr_loss if ((i == 0) and (e == 0))

else (1 - smoothing_constant) * moving_loss + (smoothing_constant) * curr_loss)

test_accuracy = evaluate_accuracy(test_data, net)

train_accuracy = evaluate_accuracy(train_data, net)

print("Epoch %s. Loss: %s, Train_acc %s, Test_acc %s" % (e, moving_loss, train_accuracy, test_accuracy))