fix mnist dataset example

Author: leewyang

examples/mnist/tf/mnist_dist_dataset.py

@@ -22,6 +22,7 @@ def map_fun(args, ctx):
   import tensorflow as tf
   import time
 
+  num_workers = args.cluster_size if args.driver_ps_nodes else args.cluster_size - args.num_ps
   worker_num = ctx.worker_num
   job_name = ctx.job_name
   task_index = ctx.task_index
@@ -43,10 +44,9 @@ def _parse_csv(ln):
     normalized_image = tf.div(image, norm)
     label_value = tf.string_to_number(lbl, tf.int32)
     label = tf.one_hot(label_value, 10)
-    return (normalized_image, label, label_value)
+    return (normalized_image, label)
 
   def _parse_tfr(example_proto):
-    print("example_proto: {}".format(example_proto))
     feature_def = {"label": tf.FixedLenFeature(10, tf.int64),
                    "image": tf.FixedLenFeature(IMAGE_PIXELS * IMAGE_PIXELS, tf.int64)}
     features = tf.parse_single_example(example_proto, feature_def)
@@ -68,10 +68,17 @@ def _parse_tfr(example_proto):
       file_pattern = os.path.join(image_dir, 'part-*')
       files = tf.gfile.Glob(file_pattern)
 
-      parse_fn = _parse_tfr if args.format == 'tfr' else _parse_csv
-      ds = tf.data.TextLineDataset(files).map(parse_fn).batch(args.batch_size)
+      if args.format == 'csv2':
+        ds = tf.data.TextLineDataset(files)
+        parse_fn = _parse_csv
+      else:  # args.format == 'tfr'
+        ds = tf.data.TFRecordDataset(files)
+        parse_fn = _parse_tfr
+
+      ds = ds.shard(num_workers, task_index).repeat(args.epochs).shuffle(args.shuffle_size)
+      ds = ds.map(parse_fn).batch(args.batch_size)
       iterator = ds.make_initializable_iterator()
-      x, y_, y_val = iterator.get_next()
+      x, y_ = iterator.get_next()
 
       # Variables of the hidden layer
       hid_w = tf.Variable(tf.truncated_normal([IMAGE_PIXELS * IMAGE_PIXELS, hidden_units],
@@ -156,8 +163,7 @@ def _parse_tfr(example_proto):
         if args.mode == "train":
           if (step % 100 == 0):
             print("{0} step: {1} accuracy: {2}".format(datetime.now().isoformat(), step, sess.run(accuracy)))
-          _, summary, step, yv = sess.run([train_op, summary_op, global_step, y_val])
-          # print("yval: {}".format(yv))
+          _, summary, step = sess.run([train_op, summary_op, global_step])
           if sv.is_chief:
             summary_writer.add_summary(summary, step)
         else:  # args.mode == "inference"

examples/mnist/tf/mnist_spark_dataset.py

@@ -21,27 +21,29 @@
 num_ps = 1
 
 parser = argparse.ArgumentParser()
-parser.add_argument("-b", "--batch_size", help="number of records per batch", type=int, default=100)
-parser.add_argument("-e", "--epochs", help="number of epochs", type=int, default=0)
-parser.add_argument("-f", "--format", help="example format: (csv2|tfr)", choices=["csv2", "tfr"], default="tfr")
-parser.add_argument("-i", "--images", help="HDFS path to MNIST images in parallelized format")
-parser.add_argument("-l", "--labels", help="HDFS path to MNIST labels in parallelized format")
-parser.add_argument("-m", "--model", help="HDFS path to save/load model during train/test", default="mnist_model")
-parser.add_argument("-n", "--cluster_size", help="number of nodes in the cluster (for Spark Standalone)", type=int, default=num_executors)
-parser.add_argument("-o", "--output", help="HDFS path to save test/inference output", default="predictions")
-parser.add_argument("-r", "--readers", help="number of reader/enqueue threads", type=int, default=1)
-parser.add_argument("-s", "--steps", help="maximum number of steps", type=int, default=1000)
-parser.add_argument("-tb", "--tensorboard", help="launch tensorboard process", action="store_true")
-parser.add_argument("-X", "--mode", help="train|inference", default="train")
-parser.add_argument("-c", "--rdma", help="use rdma connection", default=False)
-parser.add_argument("-p", "--driver_ps_nodes", help="""run tensorflow PS node on driver locally.
+parser.add_argument("--batch_size", help="number of records per batch", type=int, default=100)
+parser.add_argument("--cluster_size", help="number of nodes in the cluster (for Spark Standalone)", type=int, default=num_executors)
+parser.add_argument("--driver_ps_nodes", help="""run tensorflow PS node on driver locally.
     You will need to set cluster_size = num_executors + num_ps""", default=False)
+parser.add_argument("--epochs", help="number of epochs", type=int, default=1)
+parser.add_argument("--format", help="example format: (csv2|tfr)", choices=["csv2", "tfr"], default="tfr")
+parser.add_argument("--images", help="HDFS path to MNIST images in parallelized format")
+parser.add_argument("--labels", help="HDFS path to MNIST labels in parallelized format")
+parser.add_argument("--mode", help="train|inference", default="train")
+parser.add_argument("--model", help="HDFS path to save/load model during train/test", default="mnist_model")
+parser.add_argument("--num_ps", help="number of ps nodes", default=1)
+parser.add_argument("--output", help="HDFS path to save test/inference output", default="predictions")
+parser.add_argument("--rdma", help="use rdma connection", default=False)
+parser.add_argument("--readers", help="number of reader/enqueue threads", type=int, default=1)
+parser.add_argument("--shuffle_size", help="size of shuffle buffer", type=int, default=1000)
+parser.add_argument("--steps", help="maximum number of steps", type=int, default=1000)
+parser.add_argument("--tensorboard", help="launch tensorboard process", action="store_true")
 args = parser.parse_args()
 print("args:", args)
 
 
 print("{0} ===== Start".format(datetime.now().isoformat()))
-cluster = TFCluster.run(sc, mnist_dist_dataset.map_fun, args, args.cluster_size, num_ps, args.tensorboard,
+cluster = TFCluster.run(sc, mnist_dist_dataset.map_fun, args, args.cluster_size, args.num_ps, args.tensorboard,
                         TFCluster.InputMode.TENSORFLOW, driver_ps_nodes=args.driver_ps_nodes)
 cluster.shutdown()