MNIST: Eager execution.

Add an example showing how to train the MNIST model with eager execution
enabled.

(This change requires changes to TensorFlow made after
the 1.6 release branch was cut, i.e., will require a build
from source or TensorFlow 1.7+)

official/mnist/README.md

@@ -60,3 +60,22 @@ Result for output key probabilities:
     2.52568233e-03   4.15460236e-04]]
 ```
 
+## Experimental: Eager Execution
+
+[Eager execution](https://research.googleblog.com/2017/10/eager-execution-imperative-define-by.html)
+(an preview feature in TensorFlow 1.5) is an imperative interface to TensorFlow.
+The exact same model defined in `mnist.py` can be trained without creating a
+TensorFlow graph using:
+
+```
+python mnist_eager.py
+```
+
+## Experimental: TPU Acceleration
+
+`mnist.py` (and `mnist_eager.py`) demonstrate training a neural network to
+classify digits on CPUs and GPUs. `mnist_tpu.py` can be used to train the
+same model using TPUs for hardware acceleration. More information in
+the [tensorflow/tpu](https://github.com/tensorflow/tpu) repository.
+
+

official/mnist/mnist.py

@@ -25,8 +25,16 @@ import tensorflow as tf
 import dataset
 
 
-class Model(object):
-  """Class that defines a graph to recognize digits in the MNIST dataset."""
+class Model(tf.keras.Model):
+  """Class that defines a graph to recognize digits in the MNIST dataset.
+
+  Network structure is equivalent to:
+  https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/examples/tutorials/mnist/mnist_deep.py
+  and
+  https://github.com/tensorflow/models/blob/master/tutorials/image/mnist/convolutional.py
+
+  But written as a tf.keras and tf.layers APIs.
+  """
 
   def __init__(self, data_format):
     """Creates a model for classifying a hand-written digit.
@@ -37,6 +45,7 @@ class Model(object):
         typically faster on CPUs. See
         https://www.tensorflow.org/performance/performance_guide#data_formats
     """
+    super(Model, self).__init__()
     if data_format == 'channels_first':
       self._input_shape = [-1, 1, 28, 28]
     else:

official/mnist/mnist_eager.py

+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""MNIST model training with TensorFlow eager execution.
+
+See:
+https://research.googleblog.com/2017/10/eager-execution-imperative-define-by.html
+
+This program demonstrates training of the convolutional neural network model
+defined in mnist.py with eager execution enabled.
+
+If you are not interested in eager execution, you should ignore this file.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import os
+import sys
+import time
+
+import tensorflow as tf
+import tensorflow.contrib.eager as tfe
+import mnist
+import dataset
+
+FLAGS = None
+
+
+def loss(logits, labels):
+  return tf.reduce_mean(
+      tf.nn.sparse_softmax_cross_entropy_with_logits(
+          logits=logits, labels=labels))
+
+
+def compute_accuracy(logits, labels):
+  predictions = tf.argmax(logits, axis=1, output_type=tf.int64)
+  labels = tf.cast(labels, tf.int64)
+  batch_size = int(logits.shape[0])
+  return tf.reduce_sum(
+      tf.cast(tf.equal(predictions, labels), dtype=tf.float32)) / batch_size
+
+
+def train(model, optimizer, dataset, log_interval=None):
+  """Trains model on `dataset` using `optimizer`."""
+
+  global_step = tf.train.get_or_create_global_step()
+
+  start = time.time()
+  for (batch, (images, labels)) in enumerate(tfe.Iterator(dataset)):
+    with tf.contrib.summary.record_summaries_every_n_global_steps(10):
+      with tfe.GradientTape() as tape:
+        logits = model(images, training=True)
+        loss_value = loss(logits, labels)
+        tf.contrib.summary.scalar('loss', loss_value)
+        tf.contrib.summary.scalar('accuracy', compute_accuracy(logits, labels))
+      grads = tape.gradient(loss_value, model.variables)
+      optimizer.apply_gradients(
+          zip(grads, model.variables), global_step=global_step)
+      if log_interval and batch % log_interval == 0:
+        rate = log_interval / (time.time() - start)
+        print('Step #%d\tLoss: %.6f (%d steps/sec)' % (batch, loss_value, rate))
+        start = time.time()
+
+
+def test(model, dataset):
+  """Perform an evaluation of `model` on the examples from `dataset`."""
+  avg_loss = tfe.metrics.Mean('loss')
+  accuracy = tfe.metrics.Accuracy('accuracy')
+
+  for (images, labels) in tfe.Iterator(dataset):
+    logits = model(images, training=False)
+    avg_loss(loss(logits, labels))
+    accuracy(
+        tf.argmax(logits, axis=1, output_type=tf.int64),
+        tf.cast(labels, tf.int64))
+  print('Test set: Average loss: %.4f, Accuracy: %4f%%\n' %
+        (avg_loss.result(), 100 * accuracy.result()))
+  with tf.contrib.summary.always_record_summaries():
+    tf.contrib.summary.scalar('loss', avg_loss.result())
+    tf.contrib.summary.scalar('accuracy', accuracy.result())
+
+
+def main(_):
+  tfe.enable_eager_execution()
+
+  (device, data_format) = ('/gpu:0', 'channels_first')
+  if FLAGS.no_gpu or tfe.num_gpus() <= 0:
+    (device, data_format) = ('/cpu:0', 'channels_last')
+  print('Using device %s, and data format %s.' % (device, data_format))
+
+  # Load the datasets
+  train_ds = dataset.train(FLAGS.data_dir).shuffle(60000).batch(
+      FLAGS.batch_size)
+  test_ds = dataset.test(FLAGS.data_dir).batch(10000)
+
+  # Create the model and optimizer
+  model = mnist.Model(data_format)
+  optimizer = tf.train.MomentumOptimizer(FLAGS.lr, FLAGS.momentum)
+
+  if FLAGS.output_dir:
+    train_dir = os.path.join(FLAGS.output_dir, 'train')
+    test_dir = os.path.join(FLAGS.output_dir, 'eval')
+    tf.gfile.MakeDirs(FLAGS.output_dir)
+  else:
+    train_dir = None
+    test_dir = None
+  summary_writer = tf.contrib.summary.create_file_writer(
+      train_dir, flush_millis=10000)
+  test_summary_writer = tf.contrib.summary.create_file_writer(
+      test_dir, flush_millis=10000, name='test')
+  checkpoint_prefix = os.path.join(FLAGS.checkpoint_dir, 'ckpt')
+
+  with tf.device(device):
+    for epoch in range(1, 11):
+      with tfe.restore_variables_on_create(
+          tf.train.latest_checkpoint(FLAGS.checkpoint_dir)):
+        global_step = tf.train.get_or_create_global_step()
+        start = time.time()
+        with summary_writer.as_default():
+          train(model, optimizer, train_ds, FLAGS.log_interval)
+        end = time.time()
+        print('\nTrain time for epoch #%d (global step %d): %f' %
+              (epoch, global_step.numpy(), end - start))
+      with test_summary_writer.as_default():
+        test(model, test_ds)
+      all_variables = (model.variables + optimizer.variables() + [global_step])
+      tfe.Saver(all_variables).save(checkpoint_prefix, global_step=global_step)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--data_dir',
+      type=str,
+      default='/tmp/tensorflow/mnist/input_data',
+      help='Directory for storing input data')
+  parser.add_argument(
+      '--batch_size',
+      type=int,
+      default=100,
+      metavar='N',
+      help='input batch size for training (default: 100)')
+  parser.add_argument(
+      '--log_interval',
+      type=int,
+      default=10,
+      metavar='N',
+      help='how many batches to wait before logging training status')
+  parser.add_argument(
+      '--output_dir',
+      type=str,
+      default=None,
+      metavar='N',
+      help='Directory to write TensorBoard summaries')
+  parser.add_argument(
+      '--checkpoint_dir',
+      type=str,
+      default='/tmp/tensorflow/mnist/checkpoints/',
+      metavar='N',
+      help='Directory to save checkpoints in (once per epoch)')
+  parser.add_argument(
+      '--lr',
+      type=float,
+      default=0.01,
+      metavar='LR',
+      help='learning rate (default: 0.01)')
+  parser.add_argument(
+      '--momentum',
+      type=float,
+      default=0.5,
+      metavar='M',
+      help='SGD momentum (default: 0.5)')
+  parser.add_argument(
+      '--no_gpu',
+      action='store_true',
+      default=False,
+      help='disables GPU usage even if a GPU is available')
+
+  FLAGS, unparsed = parser.parse_known_args()
+  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)

official/mnist/mnist_eager_test.py

+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import tensorflow.contrib.eager as tfe
+
+import mnist
+import mnist_eager
+
+
+def device():
+  return "/device:GPU:0" if tfe.num_gpus() else "/device:CPU:0"
+
+
+def data_format():
+  return "channels_first" if tfe.num_gpus() else "channels_last"
+
+
+def random_dataset():
+  batch_size = 64
+  images = tf.random_normal([batch_size, 784])
+  labels = tf.random_uniform([batch_size], minval=0, maxval=10, dtype=tf.int32)
+  return tf.data.Dataset.from_tensors((images, labels))
+
+
+def train(defun=False):
+  model = mnist.Model(data_format())
+  if defun:
+    model.call = tfe.defun(model.call)
+  optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)
+  dataset = random_dataset()
+  with tf.device(device()):
+    mnist_eager.train(model, optimizer, dataset)
+
+
+def evaluate(defun=False):
+  model = mnist.Model(data_format())
+  dataset = random_dataset()
+  if defun:
+    model.call = tfe.defun(model.call)
+  with tf.device(device()):
+    mnist_eager.test(model, dataset)
+
+
+class MNISTTest(tf.test.TestCase):
+
+  def test_train(self):
+    train(defun=False)
+
+  def test_evaluate(self):
+    evaluate(defun=False)
+
+  def test_train_with_defun(self):
+    train(defun=True)
+
+  def test_evaluate_with_defun(self):
+    evaluate(defun=True)
+
+
+if __name__ == "__main__":
+  tfe.enable_eager_execution()
+  tf.test.main()