#!/usr/bin/python

import time

from tensorflow.examples.tutorials.mnist import input_data
import tensorflow as tf

import numpy as np

from simple_lenet import LeNetClassifier


def run_mnist_test():
  """ Runs the a test that trains a CNN to classify MNIST digits.
  Returns:
    A tuple containing the total elapsed time, and the average number of
    training iterations per second. """
  mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

  conv1 = LeNetClassifier.ConvLayer(kernel_width=3, kernel_height=3,
                                    feature_maps=1)
  conv2 = LeNetClassifier.ConvLayer(kernel_width=3, kernel_height=3,
                                    feature_maps=32)
  conv3 = LeNetClassifier.ConvLayer(kernel_width=3, kernel_height=3,
                                    feature_maps=64)
  network = LeNetClassifier((28, 28, 1), [conv1, conv2, conv3],
                            [4 * 4 * 128, 625], 10, batch_size=128)

  saver = tf.train.Saver()

  sess = tf.Session()
  init = tf.initialize_all_variables()
  sess.run(init)

  writer = tf.train.SummaryWriter("mnist_logs", sess.graph_def)

  print("Tensorflow: Starting MNIST test...")

  accuracy = 0
  start_time = time.time()
  iterations = 0
  while iterations < 2000:
    batch = mnist.train.next_batch(128)
    if iterations % 500 == 0:
      result = sess.run(network.predict(),
          feed_dict={network.inputs(): batch[0],
                    network.expected_outputs(): batch[1]})
      argmax = np.argmax(batch[1], axis=1)
      accuracy = np.mean(argmax == result)
      print("Tensorflow: step %d, training accuracy %s" % \
            (iterations, accuracy))

    sess.run(network.train(), feed_dict={network.inputs(): batch[0],
                                        network.expected_outputs(): batch[1]})
    iterations += 1

  # Save the network at the end.
  saver.save(sess, "Variables/test.ckpt")

  elapsed = time.time() - start_time
  speed = iterations / elapsed
  print("Tensorflow: Ran %d training iterations. (%f iter/s)" % \
      (iterations, speed))
  print("Tensorflow: MNIST test completed in %f seconds." % (elapsed))
  return (elapsed, speed)

def main():
  run_mnist_test()


if __name__ == "__main__":
  main()

@@ -12,16 +12,28 @@
 python%0A%0A
+import json%0A
 import t
@@ -1269,16 +1269,53 @@
 00:%0A    
+if iterations %25 500 == 0:%0A      test_
 batch = 
@@ -1321,20 +1321,19 @@
  mnist.t
-rain
+est
 .next_ba
@@ -1345,38 +1345,8 @@
 28)%0A
-    if iterations %25 500 == 0:%0A
     
@@ -1415,32 +1415,37 @@
 twork.inputs(): 
+test_
 batch%5B0%5D,%0A      
@@ -1478,32 +1478,37 @@
 cted_outputs(): 
+test_
 batch%5B1%5D%7D)%0A     
@@ -1527,16 +1527,21 @@
 .argmax(
+test_
 batch%5B1%5D
@@ -1628,20 +1628,19 @@
 ep %25d, t
-rain
+est
 ing accu
@@ -1689,16 +1689,56 @@
 racy))%0A%0A
+    batch = mnist.train.next_batch(128)%0A
     sess
@@ -1933,16 +1933,17 @@
  end.%0A  
+#
 saver.sa
@@ -2248,31 +2248,153 @@
 in():%0A  
-run_mnist_test(
+elapsed, speed = run_mnist_test()%0A  results = %7B%22mnist%22: %7B%22elapsed%22: elapsed, %22speed%22: speed%7D%7D%0A  print %22results=%25s%22 %25 (json.dumps(results)
 )%0A%0A%0Aif _

Make run_tests.py use testing batches.