text: # pylint: skip-file
  from __future__ import print_function
  
  import argparse
  import logging
  logging.basicConfig(level=logging.DEBUG)
  
  import numpy as np
  import mxnet as mx
  from mxnet import gluon, autograd
  from mxnet.gluon import nn
  
  from data import mnist_iterator
  
  
  # Parse CLI arguments
  
  parser = argparse.ArgumentParser(description='MXNet Gluon MNIST Example')
  parser.add_argument('--batch-size', type=int, default=100,
                      help='batch size for training and testing (default: 100)')
  parser.add_argument('--epochs', type=int, default=10,
                      help='number of epochs to train (default: 10)')
  parser.add_argument('--lr', type=float, default=0.1,
                      help='learning rate (default: 0.1)')
  parser.add_argument('--momentum', type=float, default=0.9,
                      help='SGD momentum (default: 0.9)')
  parser.add_argument('--cuda', action='store_true', default=False,
                      help='Train on GPU with CUDA')
  parser.add_argument('--log-interval', type=int, default=100, metavar='N',
                      help='how many batches to wait before logging training status')
  opt = parser.parse_args()
  
  
  # define network
  
  net = nn.Sequential()
  with net.name_scope():
      net.add(nn.Dense(128, activation='relu'))
      net.add(nn.Dense(64, activation='relu'))
      net.add(nn.Dense(10))
  
  # data
  
  train_data, val_data = mnist_iterator(batch_size=opt.batch_size, input_shape=(28*28,))
  
  # train
  
  def test(ctx):
      metric = mx.metric.Accuracy()
      val_data.reset()
      for batch in val_data:
          data = batch.data[0].as_in_context(ctx)
          label = batch.label[0].as_in_context(ctx)
          output = net(data)
          metric.update([label], [output])
  
      return metric.get()
  
  
  def train(epochs, ctx):
      # Collect all parameters from net and its children, then initialize them.
      net.collect_params().initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)
      # Trainer is for updating parameters with gradient.
      trainer = gluon.Trainer(net.collect_params(), 'sgd',
                            {'learning_rate': opt.lr, 'momentum': opt.momentum})
      metric = mx.metric.Accuracy()
      loss = gluon.loss.SoftmaxCrossEntropyLoss()
  
      for epoch in range(epochs):
          # reset data iterator and metric at begining of epoch.
          train_data.reset()
          metric.reset()
          for i, batch in enumerate(train_data):
              # Copy data to ctx if necessary
              data = batch.data[0].as_in_context(ctx)
              label = batch.label[0].as_in_context(ctx)
              # Start recording computation graph with record() section.
              # Recorded graphs can then be differentiated with backward.
              with autograd.record():
                  output = net(data)
                  L = loss(output, label)
                  L.backward()
              # take a gradient step with batch_size equal to data.shape[0]
              trainer.step(data.shape[0])
              # update metric at last.
              metric.update([label], [output])
  
              if i % opt.log_interval == 0 and i > 0:
                  name, acc = metric.get()
                  print('[Epoch %d Batch %d] Training: %s=%f'%(epoch, i, name, acc))
  
          name, acc = metric.get()
          print('[Epoch %d] Training: %s=%f'%(epoch, name, acc))
  
          name, val_acc = test(ctx)
          print('[Epoch %d] Validation: %s=%f'%(epoch, name, val_acc))
  
      net.collect_params().save('mnist.params')
  
  
  if __name__ == '__main__':
      if opt.cuda:
          ctx = mx.gpu(0)
      else:
          ctx = mx.cpu()
      train(opt.epochs, ctx)
  