Move official/mnist folder to official/r1/mnist.

PiperOrigin-RevId: 278649613

official/README.md

@@ -84,13 +84,14 @@ installable Official Models package. This is being tracked in
 
 ### Computer Vision
 
+*   [mnist](vision/image_classification): A basic model to classify digits from
+    the MNIST dataset.
 *   [resnet](vision/image_classification): A deep residual network that can be
     used to classify both CIFAR-10 and ImageNet's dataset of 1000 classes.
 *   [retinanet](vision/detection): A fast and power detector.
 
 ### Others
 
-*   [mnist](mnist): A basic model to classify digits from the MNIST dataset.
 *   [ncf](recommendation): Neural Collaborative Filtering model for
     recommendation tasks.
 

official/r1/mnist/README.md

+# MNIST in TensorFlow
+
+This directory builds a convolutional neural net to classify the [MNIST
+dataset](http://yann.lecun.com/exdb/mnist/) using the
+[tf.data](https://www.tensorflow.org/api_docs/python/tf/data),
+[tf.estimator.Estimator](https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator),
+and
+[tf.layers](https://www.tensorflow.org/api_docs/python/tf/layers)
+APIs.
+
+
+## Setup
+
+To begin, you'll simply need the latest version of TensorFlow installed.
+First make sure you've [added the models folder to your Python path]:
+
+```shell
+export PYTHONPATH="$PYTHONPATH:/path/to/models"
+```
+
+Otherwise you may encounter an error like `ImportError: No module named official.mnist`.
+
+Then to train the model, run the following:
+
+```
+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/guide/saved_model) format by using the argument `--export_dir`:
+
+```
+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/`).
+
+**Getting predictions with SavedModel**
+Use [`saved_model_cli`](https://www.tensorflow.org/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=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.
+
+The output should look similar to below:
+```
+Result for output key classes:
+[5 3]
+Result for output key probabilities:
+[[  1.53558474e-07   1.95694142e-13   1.31193523e-09   5.47467265e-03
+    5.85711526e-22   9.94520664e-01   3.48423509e-06   2.65365645e-17
+    9.78631419e-07   3.15522470e-08]
+ [  1.22413359e-04   5.87615965e-08   1.72251271e-06   9.39960718e-01
+    3.30306928e-11   2.87386645e-02   2.82353517e-02   8.21146413e-18
+    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/r1/mnist/dataset.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.
+"""tf.data.Dataset interface to the MNIST dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gzip
+import os
+import shutil
+import tempfile
+
+import numpy as np
+from six.moves import urllib
+import tensorflow as tf
+
+
+def read32(bytestream):
+  """Read 4 bytes from bytestream as an unsigned 32-bit integer."""
+  dt = np.dtype(np.uint32).newbyteorder('>')
+  return np.frombuffer(bytestream.read(4), dtype=dt)[0]
+
+
+def check_image_file_header(filename):
+  """Validate that filename corresponds to images for the MNIST dataset."""
+  with tf.io.gfile.GFile(filename, 'rb') as f:
+    magic = read32(f)
+    read32(f)  # num_images, unused
+    rows = read32(f)
+    cols = read32(f)
+    if magic != 2051:
+      raise ValueError('Invalid magic number %d in MNIST file %s' % (magic,
+                                                                     f.name))
+    if rows != 28 or cols != 28:
+      raise ValueError(
+          'Invalid MNIST file %s: Expected 28x28 images, found %dx%d' %
+          (f.name, rows, cols))
+
+
+def check_labels_file_header(filename):
+  """Validate that filename corresponds to labels for the MNIST dataset."""
+  with tf.io.gfile.GFile(filename, 'rb') as f:
+    magic = read32(f)
+    read32(f)  # num_items, unused
+    if magic != 2049:
+      raise ValueError('Invalid magic number %d in MNIST file %s' % (magic,
+                                                                     f.name))
+
+
+def download(directory, filename):
+  """Download (and unzip) a file from the MNIST dataset if not already done."""
+  filepath = os.path.join(directory, filename)
+  if tf.io.gfile.exists(filepath):
+    return filepath
+  if not tf.io.gfile.exists(directory):
+    tf.io.gfile.makedirs(directory)
+  # CVDF mirror of http://yann.lecun.com/exdb/mnist/
+  url = 'https://storage.googleapis.com/cvdf-datasets/mnist/' + filename + '.gz'
+  _, zipped_filepath = tempfile.mkstemp(suffix='.gz')
+  print('Downloading %s to %s' % (url, zipped_filepath))
+  urllib.request.urlretrieve(url, zipped_filepath)
+  with gzip.open(zipped_filepath, 'rb') as f_in, \
+      tf.io.gfile.GFile(filepath, 'wb') as f_out:
+    shutil.copyfileobj(f_in, f_out)
+  os.remove(zipped_filepath)
+  return filepath
+
+
+def dataset(directory, images_file, labels_file):
+  """Download and parse MNIST dataset."""
+
+  images_file = download(directory, images_file)
+  labels_file = download(directory, labels_file)
+
+  check_image_file_header(images_file)
+  check_labels_file_header(labels_file)
+
+  def decode_image(image):
+    # Normalize from [0, 255] to [0.0, 1.0]
+    image = tf.io.decode_raw(image, tf.uint8)
+    image = tf.cast(image, tf.float32)
+    image = tf.reshape(image, [784])
+    return image / 255.0
+
+  def decode_label(label):
+    label = tf.io.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
+    label = tf.reshape(label, [])  # label is a scalar
+    return tf.cast(label, tf.int32)
+
+  images = tf.data.FixedLengthRecordDataset(
+      images_file, 28 * 28, header_bytes=16).map(decode_image)
+  labels = tf.data.FixedLengthRecordDataset(
+      labels_file, 1, header_bytes=8).map(decode_label)
+  return tf.data.Dataset.zip((images, labels))
+
+
+def train(directory):
+  """tf.data.Dataset object for MNIST training data."""
+  return dataset(directory, 'train-images-idx3-ubyte',
+                 'train-labels-idx1-ubyte')
+
+
+def test(directory):
+  """tf.data.Dataset object for MNIST test data."""
+  return dataset(directory, 't10k-images-idx3-ubyte', 't10k-labels-idx1-ubyte')

official/r1/mnist/mnist.py

+#  Copyright 2017 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.
+"""Convolutional Neural Network Estimator for MNIST, built with tf.layers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import app as absl_app
+from absl import flags
+from six.moves import range
+import tensorflow as tf  # pylint: disable=g-bad-import-order
+
+from official.r1.mnist import dataset
+from official.utils.flags import core as flags_core
+from official.utils.logs import hooks_helper
+from official.utils.misc import distribution_utils
+from official.utils.misc import model_helpers
+
+
+LEARNING_RATE = 1e-4
+
+
+def create_model(data_format):
+  """Model 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 uses the tf.keras API.
+
+  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
+
+  Returns:
+    A tf.keras.Model.
+  """
+  if data_format == 'channels_first':
+    input_shape = [1, 28, 28]
+  else:
+    assert data_format == 'channels_last'
+    input_shape = [28, 28, 1]
+
+  l = tf.keras.layers
+  max_pool = l.MaxPooling2D(
+      (2, 2), (2, 2), padding='same', data_format=data_format)
+  # The model consists of a sequential chain of layers, so tf.keras.Sequential
+  # (a subclass of tf.keras.Model) makes for a compact description.
+  return tf.keras.Sequential(
+      [
+          l.Reshape(
+              target_shape=input_shape,
+              input_shape=(28 * 28,)),
+          l.Conv2D(
+              32,
+              5,
+              padding='same',
+              data_format=data_format,
+              activation=tf.nn.relu),
+          max_pool,
+          l.Conv2D(
+              64,
+              5,
+              padding='same',
+              data_format=data_format,
+              activation=tf.nn.relu),
+          max_pool,
+          l.Flatten(),
+          l.Dense(1024, activation=tf.nn.relu),
+          l.Dropout(0.4),
+          l.Dense(10)
+      ])
+
+
+def define_mnist_flags():
+  """Defines flags for mnist."""
+  flags_core.define_base(clean=True, train_epochs=True,
+                         epochs_between_evals=True, stop_threshold=True,
+                         num_gpu=True, hooks=True, export_dir=True,
+                         distribution_strategy=True)
+  flags_core.define_performance(inter_op=True, intra_op=True,
+                                num_parallel_calls=False,
+                                all_reduce_alg=True)
+  flags_core.define_image()
+  flags.adopt_module_key_flags(flags_core)
+  flags_core.set_defaults(data_dir='/tmp/mnist_data',
+                          model_dir='/tmp/mnist_model',
+                          batch_size=100,
+                          train_epochs=40)
+
+
+def model_fn(features, labels, mode, params):
+  """The model_fn argument for creating an Estimator."""
+  model = create_model(params['data_format'])
+  image = features
+  if isinstance(image, dict):
+    image = features['image']
+
+  if mode == tf.estimator.ModeKeys.PREDICT:
+    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.compat.v1.train.AdamOptimizer(learning_rate=LEARNING_RATE)
+
+    logits = model(image, training=True)
+    loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(labels=labels,
+                                                            logits=logits)
+    accuracy = tf.compat.v1.metrics.accuracy(
+        labels=labels, predictions=tf.argmax(logits, axis=1))
+
+    # Name tensors to be logged with LoggingTensorHook.
+    tf.identity(LEARNING_RATE, 'learning_rate')
+    tf.identity(loss, 'cross_entropy')
+    tf.identity(accuracy[1], name='train_accuracy')
+
+    # Save accuracy scalar to Tensorboard output.
+    tf.summary.scalar('train_accuracy', accuracy[1])
+
+    return tf.estimator.EstimatorSpec(
+        mode=tf.estimator.ModeKeys.TRAIN,
+        loss=loss,
+        train_op=optimizer.minimize(
+            loss,
+            tf.compat.v1.train.get_or_create_global_step()))
+  if mode == tf.estimator.ModeKeys.EVAL:
+    logits = model(image, training=False)
+    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
+    return tf.estimator.EstimatorSpec(
+        mode=tf.estimator.ModeKeys.EVAL,
+        loss=loss,
+        eval_metric_ops={
+            'accuracy':
+                tf.metrics.accuracy(
+                    labels=labels, predictions=tf.argmax(logits, axis=1)),
+        })
+
+
+def run_mnist(flags_obj):
+  """Run MNIST training and eval loop.
+
+  Args:
+    flags_obj: An object containing parsed flag values.
+  """
+  model_helpers.apply_clean(flags_obj)
+  model_function = model_fn
+
+  session_config = tf.compat.v1.ConfigProto(
+      inter_op_parallelism_threads=flags_obj.inter_op_parallelism_threads,
+      intra_op_parallelism_threads=flags_obj.intra_op_parallelism_threads,
+      allow_soft_placement=True)
+
+  distribution_strategy = distribution_utils.get_distribution_strategy(
+      distribution_strategy=flags_obj.distribution_strategy,
+      num_gpus=flags_core.get_num_gpus(flags_obj),
+      all_reduce_alg=flags_obj.all_reduce_alg)
+
+  run_config = tf.estimator.RunConfig(
+      train_distribute=distribution_strategy, session_config=session_config)
+
+  data_format = flags_obj.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=model_function,
+      model_dir=flags_obj.model_dir,
+      config=run_config,
+      params={
+          'data_format': data_format,
+      })
+
+  # Set up training and evaluation input functions.
+  def train_input_fn():
+    """Prepare data for training."""
+
+    # When choosing shuffle buffer sizes, larger sizes result in better
+    # randomness, while smaller sizes use less memory. MNIST is a small
+    # enough dataset that we can easily shuffle the full epoch.
+    ds = dataset.train(flags_obj.data_dir)
+    ds = ds.cache().shuffle(buffer_size=50000).batch(flags_obj.batch_size)
+
+    # Iterate through the dataset a set number (`epochs_between_evals`) of times
+    # during each training session.
+    ds = ds.repeat(flags_obj.epochs_between_evals)
+    return ds
+
+  def eval_input_fn():
+    return dataset.test(flags_obj.data_dir).batch(
+        flags_obj.batch_size).make_one_shot_iterator().get_next()
+
+  # Set up hook that outputs training logs every 100 steps.
+  train_hooks = hooks_helper.get_train_hooks(
+      flags_obj.hooks, model_dir=flags_obj.model_dir,
+      batch_size=flags_obj.batch_size)
+
+  # Train and evaluate model.
+  for _ in range(flags_obj.train_epochs // flags_obj.epochs_between_evals):
+    mnist_classifier.train(input_fn=train_input_fn, hooks=train_hooks)
+    eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
+    print('\nEvaluation results:\n\t%s\n' % eval_results)
+
+    if model_helpers.past_stop_threshold(flags_obj.stop_threshold,
+                                         eval_results['accuracy']):
+      break
+
+  # Export the model
+  if flags_obj.export_dir is not None:
+    image = tf.compat.v1.placeholder(tf.float32, [None, 28, 28])
+    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
+        'image': image,
+    })
+    mnist_classifier.export_savedmodel(flags_obj.export_dir, input_fn,
+                                       strip_default_attrs=True)
+
+
+def main(_):
+  run_mnist(flags.FLAGS)
+
+
+if __name__ == '__main__':
+  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
+  define_mnist_flags()
+  absl_app.run(main)

official/r1/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 os
+import time
+
+# pylint: disable=g-bad-import-order
+from absl import app as absl_app
+from absl import flags
+import tensorflow as tf
+from tensorflow.python import eager as tfe
+# pylint: enable=g-bad-import-order
+
+from official.r1.mnist import dataset as mnist_dataset
+from official.r1.mnist import mnist
+from official.utils.flags import core as flags_core
+from official.utils.misc import model_helpers
+
+
+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, step_counter, log_interval=None):
+  """Trains model on `dataset` using `optimizer`."""
+
+  start = time.time()
+  for (batch, (images, labels)) in enumerate(dataset):
+    with tf.contrib.summary.record_summaries_every_n_global_steps(
+        10, global_step=step_counter):
+      # Record the operations used to compute the loss given the input,
+      # so that the gradient of the loss with respect to the variables
+      # can be computed.
+      with tf.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=step_counter)
+      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 = tf.keras.metrics.Mean('loss', dtype=tf.float32)
+  accuracy = tf.keras.metrics.Accuracy('accuracy', dtype=tf.float32)
+
+  for (images, labels) in dataset:
+    logits = model(images, training=False)
+    avg_loss.update_state(loss(logits, labels))
+    accuracy.update_state(
+        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 run_mnist_eager(flags_obj):
+  """Run MNIST training and eval loop in eager mode.
+
+  Args:
+    flags_obj: An object containing parsed flag values.
+  """
+  tf.enable_eager_execution()
+  model_helpers.apply_clean(flags.FLAGS)
+
+  # Automatically determine device and data_format
+  (device, data_format) = ('/gpu:0', 'channels_first')
+  if flags_obj.no_gpu or not tf.test.is_gpu_available():
+    (device, data_format) = ('/cpu:0', 'channels_last')
+  # If data_format is defined in FLAGS, overwrite automatically set value.
+  if flags_obj.data_format is not None:
+    data_format = flags_obj.data_format
+  print('Using device %s, and data format %s.' % (device, data_format))
+
+  # Load the datasets
+  train_ds = mnist_dataset.train(flags_obj.data_dir).shuffle(60000).batch(
+      flags_obj.batch_size)
+  test_ds = mnist_dataset.test(flags_obj.data_dir).batch(
+      flags_obj.batch_size)
+
+  # Create the model and optimizer
+  model = mnist.create_model(data_format)
+  optimizer = tf.train.MomentumOptimizer(flags_obj.lr, flags_obj.momentum)
+
+  # Create file writers for writing TensorBoard summaries.
+  if flags_obj.output_dir:
+    # Create directories to which summaries will be written
+    # tensorboard --logdir=<output_dir>
+    # can then be used to see the recorded summaries.
+    train_dir = os.path.join(flags_obj.output_dir, 'train')
+    test_dir = os.path.join(flags_obj.output_dir, 'eval')
+    tf.gfile.MakeDirs(flags_obj.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')
+
+  # Create and restore checkpoint (if one exists on the path)
+  checkpoint_prefix = os.path.join(flags_obj.model_dir, 'ckpt')
+  step_counter = tf.train.get_or_create_global_step()
+  checkpoint = tf.train.Checkpoint(
+      model=model, optimizer=optimizer, step_counter=step_counter)
+  # Restore variables on creation if a checkpoint exists.
+  checkpoint.restore(tf.train.latest_checkpoint(flags_obj.model_dir))
+
+  # Train and evaluate for a set number of epochs.
+  with tf.device(device):
+    for _ in range(flags_obj.train_epochs):
+      start = time.time()
+      with summary_writer.as_default():
+        train(model, optimizer, train_ds, step_counter,
+              flags_obj.log_interval)
+      end = time.time()
+      print('\nTrain time for epoch #%d (%d total steps): %f' %
+            (checkpoint.save_counter.numpy() + 1,
+             step_counter.numpy(),
+             end - start))
+      with test_summary_writer.as_default():
+        test(model, test_ds)
+      checkpoint.save(checkpoint_prefix)
+
+
+def define_mnist_eager_flags():
+  """Defined flags and defaults for MNIST in eager mode."""
+  flags_core.define_base(clean=True, train_epochs=True, export_dir=True,
+                         distribution_strategy=True)
+  flags_core.define_image()
+  flags.adopt_module_key_flags(flags_core)
+
+  flags.DEFINE_integer(
+      name='log_interval', short_name='li', default=10,
+      help=flags_core.help_wrap('batches between logging training status'))
+
+  flags.DEFINE_string(
+      name='output_dir', short_name='od', default=None,
+      help=flags_core.help_wrap('Directory to write TensorBoard summaries'))
+
+  flags.DEFINE_float(name='learning_rate', short_name='lr', default=0.01,
+                     help=flags_core.help_wrap('Learning rate.'))
+
+  flags.DEFINE_float(name='momentum', short_name='m', default=0.5,
+                     help=flags_core.help_wrap('SGD momentum.'))
+
+  flags.DEFINE_bool(name='no_gpu', short_name='nogpu', default=False,
+                    help=flags_core.help_wrap(
+                        'disables GPU usage even if a GPU is available'))
+
+  flags_core.set_defaults(
+      data_dir='/tmp/tensorflow/mnist/input_data',
+      model_dir='/tmp/tensorflow/mnist/checkpoints/',
+      batch_size=100,
+      train_epochs=10,
+  )
+
+
+def main(_):
+  run_mnist_eager(flags.FLAGS)
+
+
+if __name__ == '__main__':
+  define_mnist_eager_flags()
+  absl_app.run(main=main)

official/r1/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 unittest
+
+import tensorflow as tf  # pylint: disable=g-bad-import-order
+from tensorflow.python import eager as tfe  # pylint: disable=g-bad-import-order
+
+from official.r1.mnist import mnist
+from official.r1.mnist import mnist_eager
+from official.utils.misc import keras_utils
+
+
+def device():
+  return '/device:GPU:0' if tfe.context.num_gpus() else '/device:CPU:0'
+
+
+def data_format():
+  return 'channels_first' if tfe.context.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.create_model(data_format())
+  if defun:
+    model.call = tf.function(model.call)
+  optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)
+  dataset = random_dataset()
+  with tf.device(device()):
+    mnist_eager.train(model, optimizer, dataset,
+                      step_counter=tf.train.get_or_create_global_step())
+
+
+def evaluate(defun=False):
+  model = mnist.create_model(data_format())
+  dataset = random_dataset()
+  if defun:
+    model.call = tf.function(model.call)
+  with tf.device(device()):
+    mnist_eager.test(model, dataset)
+
+
+class MNISTTest(tf.test.TestCase):
+  """Run tests for MNIST eager loop.
+
+  MNIST eager uses contrib and will not work with TF 2.0.  All tests are
+  disabled if using TF 2.0.
+  """
+
+  def setUp(self):
+    if not keras_utils.is_v2_0():
+      tf.compat.v1.enable_v2_behavior()
+    super(MNISTTest, self).setUp()
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_train(self):
+    train(defun=False)
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_evaluate(self):
+    evaluate(defun=False)
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_train_with_defun(self):
+    train(defun=True)
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_evaluate_with_defun(self):
+    evaluate(defun=True)
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/r1/mnist/mnist_test.py

+# Copyright 2017 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 time
+import unittest
+
+import tensorflow as tf  # pylint: disable=g-bad-import-order
+
+from official.r1.mnist import mnist
+from official.utils.misc import keras_utils
+
+BATCH_SIZE = 100
+
+
+def dummy_input_fn():
+  image = tf.random.uniform([BATCH_SIZE, 784])
+  labels = tf.random.uniform([BATCH_SIZE, 1], maxval=9, dtype=tf.int32)
+  return image, labels
+
+
+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):
+  """Run tests for MNIST model.
+
+  MNIST uses contrib and will not work with TF 2.0.  All tests are disabled if
+  using TF 2.0.
+  """
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  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 _ in range(3):
+      predictions = next(predictions_generator)
+      self.assertEqual(predictions['probabilities'].shape, (10,))
+      self.assertEqual(predictions['classes'].shape, ())
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def mnist_model_fn_helper(self, mode, multi_gpu=False):
+    features, labels = dummy_input_fn()
+    image_count = features.shape[0]
+    spec = mnist.model_fn(features, labels, mode, {
+        'data_format': 'channels_last',
+        'multi_gpu': multi_gpu
+    })
+
+    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
+      self.assertAllEqual(loss.shape, ())
+      self.assertEqual(loss.dtype, tf.float32)
+
+    if mode == tf.estimator.ModeKeys.EVAL:
+      eval_metric_ops = spec.eval_metric_ops
+      self.assertAllEqual(eval_metric_ops['accuracy'][0].shape, ())
+      self.assertAllEqual(eval_metric_ops['accuracy'][1].shape, ())
+      self.assertEqual(eval_metric_ops['accuracy'][0].dtype, tf.float32)
+      self.assertEqual(eval_metric_ops['accuracy'][1].dtype, tf.float32)
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_mnist_model_fn_train_mode(self):
+    self.mnist_model_fn_helper(tf.estimator.ModeKeys.TRAIN)
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_mnist_model_fn_train_mode_multi_gpu(self):
+    self.mnist_model_fn_helper(tf.estimator.ModeKeys.TRAIN, multi_gpu=True)
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_mnist_model_fn_eval_mode(self):
+    self.mnist_model_fn_helper(tf.estimator.ModeKeys.EVAL)
+
+  @unittest.skipIf(keras_utils.is_v2_0(), 'TF 1.0 only test.')
+  def test_mnist_model_fn_predict_mode(self):
+    self.mnist_model_fn_helper(tf.estimator.ModeKeys.PREDICT)
+
+
+class Benchmarks(tf.test.Benchmark):
+  """Simple speed benchmarking for MNIST."""
+
+  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.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+  tf.test.main()

official/r1/mnist/mnist_tpu.py

+# Copyright 2017 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 using TPUs.
+
+This program demonstrates training of the convolutional neural network model
+defined in mnist.py on Google Cloud TPUs (https://cloud.google.com/tpu/).
+
+If you are not interested in TPUs, you should ignore this file.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+
+# pylint: disable=g-bad-import-order
+from absl import app as absl_app  # pylint: disable=unused-import
+import tensorflow as tf
+# pylint: enable=g-bad-import-order
+
+# For open source environment, add grandparent directory for import
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(sys.path[0]))))
+
+from official.r1.mnist import dataset  # pylint: disable=wrong-import-position
+from official.r1.mnist import mnist  # pylint: disable=wrong-import-position
+
+# Cloud TPU Cluster Resolver flags
+tf.flags.DEFINE_string(
+    "tpu", default=None,
+    help="The Cloud TPU to use for training. This should be either the name "
+    "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 "
+    "url.")
+tf.flags.DEFINE_string(
+    "tpu_zone", default=None,
+    help="[Optional] GCE zone where the Cloud TPU is located in. If not "
+    "specified, we will attempt to automatically detect the GCE project from "
+    "metadata.")
+tf.flags.DEFINE_string(
+    "gcp_project", default=None,
+    help="[Optional] Project name for the Cloud TPU-enabled project. If not "
+    "specified, we will attempt to automatically detect the GCE project from "
+    "metadata.")
+
+# Model specific parameters
+tf.flags.DEFINE_string("data_dir", "",
+                       "Path to directory containing the MNIST dataset")
+tf.flags.DEFINE_string("model_dir", None, "Estimator model_dir")
+tf.flags.DEFINE_integer("batch_size", 1024,
+                        "Mini-batch size for the training. Note that this "
+                        "is the global batch size and not the per-shard batch.")
+tf.flags.DEFINE_integer("train_steps", 1000, "Total number of training steps.")
+tf.flags.DEFINE_integer("eval_steps", 0,
+                        "Total number of evaluation steps. If `0`, evaluation "
+                        "after training is skipped.")
+tf.flags.DEFINE_float("learning_rate", 0.05, "Learning rate.")
+
+tf.flags.DEFINE_bool("use_tpu", True, "Use TPUs rather than plain CPUs")
+tf.flags.DEFINE_bool("enable_predict", True, "Do some predictions at the end")
+tf.flags.DEFINE_integer("iterations", 50,
+                        "Number of iterations per TPU training loop.")
+tf.flags.DEFINE_integer("num_shards", 8, "Number of shards (TPU chips).")
+
+FLAGS = tf.flags.FLAGS
+
+
+def metric_fn(labels, logits):
+  accuracy = tf.metrics.accuracy(
+      labels=labels, predictions=tf.argmax(logits, axis=1))
+  return {"accuracy": accuracy}
+
+
+def model_fn(features, labels, mode, params):
+  """model_fn constructs the ML model used to predict handwritten digits."""
+
+  del params
+  image = features
+  if isinstance(image, dict):
+    image = features["image"]
+
+  model = mnist.create_model("channels_last")
+
+  if mode == tf.estimator.ModeKeys.PREDICT:
+    logits = model(image, training=False)
+    predictions = {
+        'class_ids': tf.argmax(logits, axis=1),
+        'probabilities': tf.nn.softmax(logits),
+    }
+    return tf.contrib.tpu.TPUEstimatorSpec(mode, predictions=predictions)
+
+  logits = model(image, training=(mode == tf.estimator.ModeKeys.TRAIN))
+  loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
+
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    learning_rate = tf.train.exponential_decay(
+        FLAGS.learning_rate,
+        tf.train.get_global_step(),
+        decay_steps=100000,
+        decay_rate=0.96)
+    optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
+    if FLAGS.use_tpu:
+      optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)
+    return tf.contrib.tpu.TPUEstimatorSpec(
+        mode=mode,
+        loss=loss,
+        train_op=optimizer.minimize(loss, tf.train.get_global_step()))
+
+  if mode == tf.estimator.ModeKeys.EVAL:
+    return tf.contrib.tpu.TPUEstimatorSpec(
+        mode=mode, loss=loss, eval_metrics=(metric_fn, [labels, logits]))
+
+
+def train_input_fn(params):
+  """train_input_fn defines the input pipeline used for training."""
+  batch_size = params["batch_size"]
+  data_dir = params["data_dir"]
+  # Retrieves the batch size for the current shard. The # of shards is
+  # computed according to the input pipeline deployment. See
+  # `tf.contrib.tpu.RunConfig` for details.
+  ds = dataset.train(data_dir).cache().repeat().shuffle(
+      buffer_size=50000).batch(batch_size, drop_remainder=True)
+  return ds
+
+
+def eval_input_fn(params):
+  batch_size = params["batch_size"]
+  data_dir = params["data_dir"]
+  ds = dataset.test(data_dir).batch(batch_size, drop_remainder=True)
+  return ds
+
+
+def predict_input_fn(params):
+  batch_size = params["batch_size"]
+  data_dir = params["data_dir"]
+  # Take out top 10 samples from test data to make the predictions.
+  ds = dataset.test(data_dir).take(10).batch(batch_size)
+  return ds
+
+
+def main(argv):
+  del argv  # Unused.
+  tf.logging.set_verbosity(tf.logging.INFO)
+
+  tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
+      FLAGS.tpu,
+      zone=FLAGS.tpu_zone,
+      project=FLAGS.gcp_project
+  )
+
+  run_config = tf.contrib.tpu.RunConfig(
+      cluster=tpu_cluster_resolver,
+      model_dir=FLAGS.model_dir,
+      session_config=tf.ConfigProto(
+          allow_soft_placement=True, log_device_placement=True),
+      tpu_config=tf.contrib.tpu.TPUConfig(FLAGS.iterations, FLAGS.num_shards),
+  )
+
+  estimator = tf.contrib.tpu.TPUEstimator(
+      model_fn=model_fn,
+      use_tpu=FLAGS.use_tpu,
+      train_batch_size=FLAGS.batch_size,
+      eval_batch_size=FLAGS.batch_size,
+      predict_batch_size=FLAGS.batch_size,
+      params={"data_dir": FLAGS.data_dir},
+      config=run_config)
+  # TPUEstimator.train *requires* a max_steps argument.
+  estimator.train(input_fn=train_input_fn, max_steps=FLAGS.train_steps)
+  # TPUEstimator.evaluate *requires* a steps argument.
+  # Note that the number of examples used during evaluation is
+  # --eval_steps * --batch_size.
+  # So if you change --batch_size then change --eval_steps too.
+  if FLAGS.eval_steps:
+    estimator.evaluate(input_fn=eval_input_fn, steps=FLAGS.eval_steps)
+
+  # Run prediction on top few samples of test data.
+  if FLAGS.enable_predict:
+    predictions = estimator.predict(input_fn=predict_input_fn)
+
+    for pred_dict in predictions:
+      template = ('Prediction is "{}" ({:.1f}%).')
+
+      class_id = pred_dict['class_ids']
+      probability = pred_dict['probabilities'][class_id]
+
+      print(template.format(class_id, 100 * probability))
+
+
+if __name__ == "__main__":
+  absl_app.run(main)