cleaned up the optimizer and main.py a little

main.py

@@ -16,18 +16,18 @@
 """Trains and Evaluates the MNIST network using a feed dictionary."""
 # pylint: disable=missing-docstring
 
-import tensorflow as tf
-
+from neural_network_optimizer import NeuralNetworkOptimizer, timed_run
 import input_data
 from mnist_graph import MNISTGraph
 
+import tensorflow as tf
+
 # Basic model parameters as external flags.
-from neural_network_optimizer import NeuralNetworkOptimizer, timed_run
 
 flags = tf.app.flags
 FLAGS = flags.FLAGS
 flags.DEFINE_float('learning_rate', 0.01, 'Initial learning rate.')
-flags.DEFINE_float('training_precision', 0.9, 'Desired training precision.')
+flags.DEFINE_float('training_precision', 0.0, 'Precision for geometry optimization runs.')
 flags.DEFINE_float('desired_precision', 0.95, 'Desired training precision.')
 flags.DEFINE_integer('max_steps', 2000, 'Number of steps to run trainer.')
 flags.DEFINE_integer('hidden1', 128, 'Number of units in hidden layer 1.')
@@ -42,28 +42,40 @@ def run_training():
     # Get the sets of images and labels for training, validation, and test on MNIST.
     data_sets = input_data.read_data_sets(FLAGS.train_dir, FLAGS.fake_data)
 
-    optimizer = NeuralNetworkOptimizer(MNISTGraph, FLAGS, True)
-    best_geometry = optimizer.brute_force_optimal_network_geometry(data_sets, FLAGS.training_precision)
-    print(best_geometry)
+    geometry = get_network_geometry(data_sets)
 
-    graph, timing_info = timed_run(run_final_training, best_geometry, data_sets)
+    graph, cpu, wall = timed_run(run_final_training, geometry, data_sets)
 
-    print(timing_info, graph.precision, graph.step)
+    print(NeuralNetworkOptimizer.TimingInfo(cpu, wall, graph.precision, graph.step, geometry))
+
+    return graph
+
+
+def get_network_geometry(data_sets):
+    if FLAGS.training_precision:
+        optimizer = NeuralNetworkOptimizer(
+            MNISTGraph, FLAGS.training_precision, FLAGS.learning_rate, verbose=True
+        )
+        geometry = optimizer.brute_force_optimal_network_geometry(data_sets)
+        print('Best geometry found:', geometry)
+    else:
+        geometry = (FLAGS.hidden1, FLAGS.hidden2)
+    return geometry
 
 
-def run_final_training(best_geometry, data_sets):
+def run_final_training(geometry, data_sets):
     with tf.Graph().as_default():
         graph = MNISTGraph(
             learning_rate=FLAGS.learning_rate,
-            hidden1=best_geometry[3][0], hidden2=best_geometry[3][1], hidden3=best_geometry[3][2],
+            hidden1=geometry[0], hidden2=geometry[1], hidden3=geometry[2],
             batch_size=FLAGS.batch_size, train_dir=FLAGS.train_dir
         )
         graph.train(data_sets, FLAGS.max_steps, precision=FLAGS.desired_precision, steps_between_checks=250)
     return graph
 
 
 def main(_):
-    run_training()
+    graph = run_training()
 
 
 if __name__ == '__main__':

neural_network_graph.py

@@ -99,10 +99,9 @@ def evaluation(self, logits, labels):
           A scalar int32 tensor with the number of examples (out of batch_size)
           that were predicted correctly.
         """
-        # For a classifier model, we can use the in_top_k Op.
-        # It returns a bool tensor with shape [batch_size] that is true for
-        # the examples where the label's is was in the top k (here k=1)
-        # of all logits for that example.
+        # For a classifier model, we can use the in_top_k Op. It returns a bool
+        # tensor with shape [batch_size] that is true for the examples where the
+        # label's is was in the top k (here k=1) of all logits for that example.
         correct = tf.nn.in_top_k(logits, labels, 1)
         # Return the number of true entries.
         return tf.reduce_sum(tf.cast(correct, tf.int32))

neural_network_optimizer.py

@@ -7,50 +7,73 @@
 
 class NeuralNetworkOptimizer:
 
-    def __init__(self, tested_network, FLAGS, verbose=False):
+    DEFAULT_LEARNING_RATE = 0.1
+
+    class TimingInfo:
+
+        def __init__(self, cpu_time, wall_time, precision, step, layers):
+            self.cpu_time = cpu_time
+            self.wall_time = wall_time
+            self.precision = precision
+            self.step = step
+            self.layers = layers
+
+        def __str__(self):
+            return 'CPU: {:7.3f}s Wall: {:7.3f}s Precision: {:5.2f}% Iterations: {:4d} Geometry: {}'.format(
+                self.cpu_time, self.wall_time, 100.*self.precision, self.step, str(self.layers)
+            )
+
+        def __repr__(self):
+            return str(self)
+
+    def __init__(self, tested_network, training_precision, learnning_rate=None, verbose=False):
         self.tested_network = tested_network
         self.verbose = verbose
-        self.FLAGS = FLAGS
+        self.learning_rate = learnning_rate if learnning_rate else self.DEFAULT_LEARNING_RATE
+        self.training_precision = training_precision
 
-    def brute_force_optimal_network_geometry(self, data_sets, desired_precision, max_steps=10000):
+    def brute_force_optimal_network_geometry(self, data_sets, max_steps=10000):
+        results = self.time_all_tested_geometries(data_sets, max_steps)
+        return results[0].layers
+
+    def time_all_tested_geometries(self, data_sets, max_steps):
         results = []
-        for layer1_size in (32, 64, 96, 128):
-            for layer2_size in (32, 64, 96, 128):
-                if layer2_size > layer1_size:
-                    continue
-                for layer3_size in (None, 32):
-                    run_info = self.timed_run_training(
-                        data_sets,
-                        layer1_size, layer2_size, layer3_size,
-                        desired_precision=desired_precision, max_steps=max_steps
-                    )
-                    if self.verbose: print(run_info)
-                    results.append(run_info)
-
-        results = sorted(results, key=lambda r: r[0]['cpu_time'])
-        if self.verbose: pprint.pprint(results)
-        return results[0]
-
-    def timed_run_training(self, data_sets, layer1_size, layer2_size, layer3_size, desired_precision=0.9, max_steps=10000):
-        graph, timing = timed_run(self.run_training_once, data_sets, desired_precision, layer1_size, layer2_size, layer3_size, max_steps)
-        return timing, graph.precision, graph.step, (layer1_size, layer2_size, layer3_size)
-
-    def run_training_once(self, data_sets, desired_precision, layer1_size, layer2_size, layer3_size, max_steps):
+        for layer1_size, layer2_size, layer3_size in self.get_network_geometries():
+            run_info = self.timed_run_training(
+                data_sets, layer1_size, layer2_size, layer3_size, max_steps=max_steps
+            )
+            if self.verbose: print(run_info)
+            results.append(run_info)
+        results = sorted(results, key=lambda r: r.cpu_time)
+        if self.verbose: pprint.pprint(results, width=100)
+        return results
+
+    def get_network_geometries(self):
+        return ((l1, l2, l3)
+                for l1 in (32, 64, 96, 128)
+                for l2 in (32, 64, 96, 128) if l2 <= l1
+                for l3 in (None, 32))
+
+    def brute_force_optimize_learning_rate(self):
+        raise NotImplemented()
+
+    def timed_run_training(self, data_sets, layer1_size, layer2_size, layer3_size, max_steps=10000):
+        graph, cpu, wall = timed_run(self.run_training_once, data_sets, layer1_size, layer2_size, layer3_size, max_steps)
+        return self.TimingInfo(cpu, wall, graph.precision, graph.step, (layer1_size, layer2_size, layer3_size))
+
+    def run_training_once(self, data_sets, layer1_size, layer2_size, layer3_size, max_steps):
         # Tell TensorFlow that the model will be built into the default Graph.
         with tf.Graph().as_default():
             graph = self.tested_network(
-                learning_rate=self.FLAGS.learning_rate,
-                hidden1=layer1_size,
-                hidden2=layer2_size, hidden3=layer3_size,
-                batch_size=self.FLAGS.batch_size, train_dir=self.FLAGS.train_dir,
+                learning_rate=self.learning_rate,
+                hidden1=layer1_size, hidden2=layer2_size, hidden3=layer3_size,
                 verbose=False
             )
-            graph.train(data_sets, max_steps, precision=desired_precision, steps_between_checks=50)
+            graph.train(data_sets, max_steps, precision=self.training_precision, steps_between_checks=50)
         return graph
 
 
 def timed_run(function, *args, **kwargs):
         start_cpu_time, start_wall_time = time.process_time(), time.time()
         returned = function(*args, **kwargs)
-        timing_info = {'cpu_time': time.process_time()-start_cpu_time, 'wall_time': time.time()-start_wall_time}
-        return returned, timing_info
+        return returned, time.process_time()-start_cpu_time, time.time()-start_wall_time