Merge pull request #3024 from asimshankar/mnist

[mnist]: Updates

official/mnist/README.md

@@ -21,12 +21,19 @@ python mnist.py
 The model will begin training and will automatically evaluate itself on the
 validation data.
 
+Illustrative unit tests and benchmarks can be run with:
+
+```
+python mnist_test.py
+python mnist_test.py --benchmarks=.
+```
+
 ## Exporting the model
 
 You can export the model into Tensorflow [SavedModel](https://www.tensorflow.org/programmers_guide/saved_model) format by using the argument `--export_dir`:
 
 ```
-python mnist.py --export_dir /tmp/mnist_saved_model 
+python mnist.py --export_dir /tmp/mnist_saved_model
 ```
 
 The SavedModel will be saved in a timestamped directory under `/tmp/mnist_saved_model/` (e.g. `/tmp/mnist_saved_model/1513630966/`).
@@ -35,7 +42,7 @@ The SavedModel will be saved in a timestamped directory under `/tmp/mnist_saved_
 Use [`saved_model_cli`](https://www.tensorflow.org/programmers_guide/saved_model#cli_to_inspect_and_execute_savedmodel) to inspect and execute the SavedModel.
 
 ```
-saved_model_cli run --dir /tmp/mnist_saved_model/TIMESTAMP --tag_set serve --signature_def classify --inputs image_raw=examples.npy
+saved_model_cli run --dir /tmp/mnist_saved_model/TIMESTAMP --tag_set serve --signature_def classify --inputs image=examples.npy
 ```
 
 `examples.npy` contains the data from `example5.png` and `example3.png` in a numpy array, in that order. The array values are normalized to values between 0 and 1.

official/mnist/mnist.py

@@ -63,6 +63,7 @@ parser.add_argument(
     type=str,
     help='The directory where the exported SavedModel will be stored.')
 
+
 def train_dataset(data_dir):
   """Returns a tf.data.Dataset yielding (image, label) pairs for training."""
   data = input_data.read_data_sets(data_dir, one_hot=True).train
@@ -75,141 +76,115 @@ def eval_dataset(data_dir):
   return tf.data.Dataset.from_tensors((data.images, data.labels))
 
 
-def mnist_model(inputs, mode, data_format):
-  """Takes the MNIST inputs and mode and outputs a tensor of logits."""
-  # Input Layer
-  # Reshape X to 4-D tensor: [batch_size, width, height, channels]
-  # MNIST images are 28x28 pixels, and have one color channel
-  inputs = tf.reshape(inputs, [-1, 28, 28, 1])
-
-  if data_format is None:
-    # When running on GPU, transpose the data from channels_last (NHWC) to
-    # channels_first (NCHW) to improve performance.
-    # See https://www.tensorflow.org/performance/performance_guide#data_formats
-    data_format = ('channels_first'
-                   if tf.test.is_built_with_cuda() else 'channels_last')
-
-  if data_format == 'channels_first':
-    inputs = tf.transpose(inputs, [0, 3, 1, 2])
-
-  # Convolutional Layer #1
-  # Computes 32 features using a 5x5 filter with ReLU activation.
-  # Padding is added to preserve width and height.
-  # Input Tensor Shape: [batch_size, 28, 28, 1]
-  # Output Tensor Shape: [batch_size, 28, 28, 32]
-  conv1 = tf.layers.conv2d(
-      inputs=inputs,
-      filters=32,
-      kernel_size=[5, 5],
-      padding='same',
-      activation=tf.nn.relu,
-      data_format=data_format)
-
-  # Pooling Layer #1
-  # First max pooling layer with a 2x2 filter and stride of 2
-  # Input Tensor Shape: [batch_size, 28, 28, 32]
-  # Output Tensor Shape: [batch_size, 14, 14, 32]
-  pool1 = tf.layers.max_pooling2d(
-      inputs=conv1, pool_size=[2, 2], strides=2, data_format=data_format)
-
-  # Convolutional Layer #2
-  # Computes 64 features using a 5x5 filter.
-  # Padding is added to preserve width and height.
-  # Input Tensor Shape: [batch_size, 14, 14, 32]
-  # Output Tensor Shape: [batch_size, 14, 14, 64]
-  conv2 = tf.layers.conv2d(
-      inputs=pool1,
-      filters=64,
-      kernel_size=[5, 5],
-      padding='same',
-      activation=tf.nn.relu,
-      data_format=data_format)
-
-  # Pooling Layer #2
-  # Second max pooling layer with a 2x2 filter and stride of 2
-  # Input Tensor Shape: [batch_size, 14, 14, 64]
-  # Output Tensor Shape: [batch_size, 7, 7, 64]
-  pool2 = tf.layers.max_pooling2d(
-      inputs=conv2, pool_size=[2, 2], strides=2, data_format=data_format)
-
-  # Flatten tensor into a batch of vectors
-  # Input Tensor Shape: [batch_size, 7, 7, 64]
-  # Output Tensor Shape: [batch_size, 7 * 7 * 64]
-  pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
-
-  # Dense Layer
-  # Densely connected layer with 1024 neurons
-  # Input Tensor Shape: [batch_size, 7 * 7 * 64]
-  # Output Tensor Shape: [batch_size, 1024]
-  dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)
-
-  # Add dropout operation; 0.6 probability that element will be kept
-  dropout = tf.layers.dropout(
-      inputs=dense, rate=0.4, training=(mode == tf.estimator.ModeKeys.TRAIN))
-
-  # Logits layer
-  # Input Tensor Shape: [batch_size, 1024]
-  # Output Tensor Shape: [batch_size, 10]
-  logits = tf.layers.dense(inputs=dropout, units=10)
-  return logits
-
-
-def mnist_model_fn(features, labels, mode, params):
-  """Model function for MNIST."""
-  if mode == tf.estimator.ModeKeys.PREDICT and isinstance(features,dict):
-    features = features['image_raw']
-  
-  logits = mnist_model(features, mode, params['data_format'])
-
-  predictions = {
-      'classes': tf.argmax(input=logits, axis=1),
-      'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
-  }
+class Model(object):
+  """Class that defines a graph to recognize digits in the MNIST dataset."""
+
+  def __init__(self, data_format):
+    """Creates a model for classifying a hand-written digit.
+
+    Args:
+      data_format: Either 'channels_first' or 'channels_last'.
+        'channels_first' is typically faster on GPUs while 'channels_last' is
+        typically faster on CPUs. See
+        https://www.tensorflow.org/performance/performance_guide#data_formats
+    """
+    if data_format == 'channels_first':
+      self._input_shape = [-1, 1, 28, 28]
+    else:
+      assert data_format == 'channels_last'
+      self._input_shape = [-1, 28, 28, 1]
+
+    self.conv1 = tf.layers.Conv2D(
+        32, 5, padding='same', data_format=data_format, activation=tf.nn.relu)
+    self.conv2 = tf.layers.Conv2D(
+        64, 5, padding='same', data_format=data_format, activation=tf.nn.relu)
+    self.fc1 = tf.layers.Dense(1024, activation=tf.nn.relu)
+    self.fc2 = tf.layers.Dense(10)
+    self.dropout = tf.layers.Dropout(0.4)
+    self.max_pool2d = tf.layers.MaxPooling2D(
+        (2, 2), (2, 2), padding='same', data_format=data_format)
+
+  def __call__(self, inputs, training):
+    """Add operations to classify a batch of input images.
+
+    Args:
+      inputs: A Tensor representing a batch of input images.
+      training: A boolean. Set to True to add operations required only when
+        training the classifier.
+
+    Returns:
+      A logits Tensor with shape [<batch_size>, 10].
+    """
+    y = tf.reshape(inputs, self._input_shape)
+    y = self.conv1(y)
+    y = self.max_pool2d(y)
+    y = self.conv2(y)
+    y = self.max_pool2d(y)
+    y = tf.layers.flatten(y)
+    y = self.fc1(y)
+    y = self.dropout(y, training=training)
+    return self.fc2(y)
+
+
+def model_fn(features, labels, mode, params):
+  """The model_fn argument for creating an Estimator."""
+  model = Model(params['data_format'])
+  image = features
+  if isinstance(image, dict):
+    image = features['image']
 
   if mode == tf.estimator.ModeKeys.PREDICT:
-    export_outputs={'classify': tf.estimator.export.PredictOutput(predictions)}
-    return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions,
-                                      export_outputs=export_outputs)
-
-  loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)
-
-  # Configure the training op
+    logits = model(image, training=False)
+    predictions = {
+        'classes': tf.argmax(logits, axis=1),
+        'probabilities': tf.nn.softmax(logits),
+    }
+    return tf.estimator.EstimatorSpec(
+        mode=tf.estimator.ModeKeys.PREDICT,
+        predictions=predictions,
+        export_outputs={
+            'classify': tf.estimator.export.PredictOutput(predictions)
+        })
   if mode == tf.estimator.ModeKeys.TRAIN:
     optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
-    train_op = optimizer.minimize(loss, tf.train.get_or_create_global_step())
-  else:
-    train_op = None
-
-  accuracy = tf.metrics.accuracy(
-      tf.argmax(labels, axis=1), predictions['classes'])
-  metrics = {'accuracy': accuracy}
-
-  # Create a tensor named train_accuracy for logging purposes
-  tf.identity(accuracy[1], name='train_accuracy')
-  tf.summary.scalar('train_accuracy', accuracy[1])
-
-  return tf.estimator.EstimatorSpec(
-      mode=mode,
-      predictions=predictions,
-      loss=loss,
-      train_op=train_op,
-      eval_metric_ops=metrics)
+    logits = model(image, training=True)
+    loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)
+    accuracy = tf.metrics.accuracy(
+        labels=tf.argmax(labels, axis=1), predictions=tf.argmax(logits, axis=1))
+    # Name the accuracy tensor 'train_accuracy' to demonstrate the
+    # LoggingTensorHook.
+    tf.identity(accuracy[1], name='train_accuracy')
+    tf.summary.scalar('train_accuracy', accuracy[1])
+    return tf.estimator.EstimatorSpec(
+        mode=tf.estimator.ModeKeys.TRAIN,
+        loss=loss,
+        train_op=optimizer.minimize(loss, tf.train.get_or_create_global_step()))
+  if mode == tf.estimator.ModeKeys.EVAL:
+    logits = model(image, training=False)
+    loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)
+    return tf.estimator.EstimatorSpec(
+        mode=tf.estimator.ModeKeys.EVAL,
+        loss=loss,
+        eval_metric_ops={
+            'accuracy':
+                tf.metrics.accuracy(
+                    labels=tf.argmax(labels, axis=1),
+                    predictions=tf.argmax(logits, axis=1)),
+        })
 
 
 def main(unused_argv):
-  # Create the Estimator
+  data_format = FLAGS.data_format
+  if data_format is None:
+    data_format = ('channels_first'
+                   if tf.test.is_built_with_cuda() else 'channels_last')
   mnist_classifier = tf.estimator.Estimator(
-      model_fn=mnist_model_fn,
+      model_fn=model_fn,
       model_dir=FLAGS.model_dir,
       params={
-          'data_format': FLAGS.data_format
+          'data_format': data_format
       })
 
-  # Set up training hook that logs the training accuracy every 100 steps.
-  tensors_to_log = {'train_accuracy': 'train_accuracy'}
-  logging_hook = tf.train.LoggingTensorHook(
-      tensors=tensors_to_log, every_n_iter=100)
-
   # Train the model
   def train_input_fn():
     # When choosing shuffle buffer sizes, larger sizes result in better
@@ -221,6 +196,10 @@ def main(unused_argv):
     (images, labels) = dataset.make_one_shot_iterator().get_next()
     return (images, labels)
 
+  # Set up training hook that logs the training accuracy every 100 steps.
+  tensors_to_log = {'train_accuracy': 'train_accuracy'}
+  logging_hook = tf.train.LoggingTensorHook(
+      tensors=tensors_to_log, every_n_iter=100)
   mnist_classifier.train(input_fn=train_input_fn, hooks=[logging_hook])
 
   # Evaluate the model and print results
@@ -233,10 +212,11 @@ def main(unused_argv):
 
   # Export the model
   if FLAGS.export_dir is not None:
-    image = tf.placeholder(tf.float32,[None, 28, 28])
-    serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(
-        {"image_raw":image})
-    mnist_classifier.export_savedmodel(FLAGS.export_dir, serving_input_fn)
+    image = tf.placeholder(tf.float32, [None, 28, 28])
+    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
+        'image': image,
+    })
+    mnist_classifier.export_savedmodel(FLAGS.export_dir, input_fn)
 
 
 if __name__ == '__main__':

official/mnist/mnist_test.py

@@ -18,30 +18,63 @@ from __future__ import division
 from __future__ import print_function
 
 import tensorflow as tf
+import time
 
 import mnist
 
-tf.logging.set_verbosity(tf.logging.ERROR)
+BATCH_SIZE = 100
 
 
-class BaseTest(tf.test.TestCase):
+def dummy_input_fn():
+  image = tf.random_uniform([BATCH_SIZE, 784])
+  labels = tf.random_uniform([BATCH_SIZE], maxval=9, dtype=tf.int32)
+  return image, tf.one_hot(labels, 10)
 
-  def input_fn(self):
-    features = tf.random_uniform([55000, 784])
-    labels = tf.random_uniform([55000], maxval=9, dtype=tf.int32)
-    return features, tf.one_hot(labels, 10)
+
+def make_estimator():
+  data_format = 'channels_last'
+  if tf.test.is_built_with_cuda():
+    data_format = 'channels_first'
+  return tf.estimator.Estimator(
+      model_fn=mnist.model_fn, params={
+          'data_format': data_format
+      })
+
+
+class Tests(tf.test.TestCase):
+
+  def test_mnist(self):
+    classifier = make_estimator()
+    classifier.train(input_fn=dummy_input_fn, steps=2)
+    eval_results = classifier.evaluate(input_fn=dummy_input_fn, steps=1)
+
+    loss = eval_results['loss']
+    global_step = eval_results['global_step']
+    accuracy = eval_results['accuracy']
+    self.assertEqual(loss.shape, ())
+    self.assertEqual(2, global_step)
+    self.assertEqual(accuracy.shape, ())
+
+    input_fn = lambda: tf.random_uniform([3, 784])
+    predictions_generator = classifier.predict(input_fn)
+    for i in range(3):
+      predictions = next(predictions_generator)
+      self.assertEqual(predictions['probabilities'].shape, (10,))
+      self.assertEqual(predictions['classes'].shape, ())
 
   def mnist_model_fn_helper(self, mode):
-    features, labels = self.input_fn()
+    features, labels = dummy_input_fn()
     image_count = features.shape[0]
-    spec = mnist.mnist_model_fn(
-        features, labels, mode, {'data_format': 'channels_last'})
+    spec = mnist.model_fn(features, labels, mode, {
+        'data_format': 'channels_last'
+    })
 
-    predictions = spec.predictions
-    self.assertAllEqual(predictions['probabilities'].shape, (image_count, 10))
-    self.assertEqual(predictions['probabilities'].dtype, tf.float32)
-    self.assertAllEqual(predictions['classes'].shape, (image_count,))
-    self.assertEqual(predictions['classes'].dtype, tf.int64)
+    if mode == tf.estimator.ModeKeys.PREDICT:
+      predictions = spec.predictions
+      self.assertAllEqual(predictions['probabilities'].shape, (image_count, 10))
+      self.assertEqual(predictions['probabilities'].dtype, tf.float32)
+      self.assertAllEqual(predictions['classes'].shape, (image_count,))
+      self.assertEqual(predictions['classes'].dtype, tf.int64)
 
     if mode != tf.estimator.ModeKeys.PREDICT:
       loss = spec.loss
@@ -65,5 +98,31 @@ class BaseTest(tf.test.TestCase):
     self.mnist_model_fn_helper(tf.estimator.ModeKeys.PREDICT)
 
 
+class Benchmarks(tf.test.Benchmark):
+
+  def benchmark_train_step_time(self):
+    classifier = make_estimator()
+    # Run one step to warmup any use of the GPU.
+    classifier.train(input_fn=dummy_input_fn, steps=1)
+
+    have_gpu = tf.test.is_gpu_available()
+    num_steps = 1000 if have_gpu else 100
+    name = 'train_step_time_%s' % ('gpu' if have_gpu else 'cpu')
+
+    start = time.time()
+    classifier.train(input_fn=dummy_input_fn, steps=num_steps)
+    end = time.time()
+
+    wall_time = (end - start) / num_steps
+    self.report_benchmark(
+        iters=num_steps,
+        wall_time=wall_time,
+        name=name,
+        extras={
+            'examples_per_sec': BATCH_SIZE / wall_time
+        })
+
+
 if __name__ == '__main__':
+  tf.logging.set_verbosity(tf.logging.ERROR)
   tf.test.main()