Merge pull request #3034 from asimshankar/mnist-flags

[mnist]: Make flags function scoped (not module scoped).

official/mnist/mnist.py

@@ -24,45 +24,6 @@ import sys
 import tensorflow as tf
 from tensorflow.examples.tutorials.mnist import input_data
 
-parser = argparse.ArgumentParser()
-
-# Basic model parameters.
-parser.add_argument(
-    '--batch_size',
-    type=int,
-    default=100,
-    help='Number of images to process in a batch')
-
-parser.add_argument(
-    '--data_dir',
-    type=str,
-    default='/tmp/mnist_data',
-    help='Path to directory containing the MNIST dataset')
-
-parser.add_argument(
-    '--model_dir',
-    type=str,
-    default='/tmp/mnist_model',
-    help='The directory where the model will be stored.')
-
-parser.add_argument(
-    '--train_epochs', type=int, default=40, help='Number of epochs to train.')
-
-parser.add_argument(
-    '--data_format',
-    type=str,
-    default=None,
-    choices=['channels_first', 'channels_last'],
-    help='A flag to override the data format used in the model. channels_first '
-    'provides a performance boost on GPU but is not always compatible '
-    'with CPU. If left unspecified, the data format will be chosen '
-    'automatically based on whether TensorFlow was built for CPU or GPU.')
-
-parser.add_argument(
-    '--export_dir',
-    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."""
@@ -220,6 +181,38 @@ def main(unused_argv):
 
 
 if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--batch_size',
+      type=int,
+      default=100,
+      help='Number of images to process in a batch')
+  parser.add_argument(
+      '--data_dir',
+      type=str,
+      default='/tmp/mnist_data',
+      help='Path to directory containing the MNIST dataset')
+  parser.add_argument(
+      '--model_dir',
+      type=str,
+      default='/tmp/mnist_model',
+      help='The directory where the model will be stored.')
+  parser.add_argument(
+      '--train_epochs', type=int, default=40, help='Number of epochs to train.')
+  parser.add_argument(
+      '--data_format',
+      type=str,
+      default=None,
+      choices=['channels_first', 'channels_last'],
+      help='A flag to override the data format used in the model. channels_first '
+      'provides a performance boost on GPU but is not always compatible '
+      'with CPU. If left unspecified, the data format will be chosen '
+      'automatically based on whether TensorFlow was built for CPU or GPU.')
+  parser.add_argument(
+      '--export_dir',
+      type=str,
+      help='The directory where the exported SavedModel will be stored.')
+
   tf.logging.set_verbosity(tf.logging.INFO)
   FLAGS, unparsed = parser.parse_known_args()
   tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)