mnist.py

from __future__ import print_function

import math
import os

import tensorflow as tf
import clusterone
from clusterone import get_data_path, get_logs_path

from tensorflow.examples.tutorials.mnist import mnist
from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets

PATH_TO_LOCAL_LOGS = os.path.abspath(os.path.expanduser('~/Documents/mnist/logs'))
ROOT_PATH_TO_LOCAL_DATA = os.path.abspath(os.path.expanduser('~/Documents/data/mnist'))

try:
    job_name = os.environ['JOB_NAME']
    task_index = os.environ['TASK_INDEX']
    ps_hosts = os.environ['PS_HOSTS']
    worker_hosts = os.environ['WORKER_HOSTS']
except:
    job_name = None
    task_index = 0
    ps_hosts = None
    worker_hosts = None

flags = tf.app.flags

for param in os.environ.keys():
    print("%s: %s " % (param, os.environ[param]))

try:
    print("Clusterne CLI Version: %s" % clusterone.__version)
except Exception as e:
    print('Cannot get CLI version')

# Flags for configuring the distributed task
flags.DEFINE_string("job_name", job_name,
                    "job name: worker or ps")
flags.DEFINE_integer("task_index", task_index,
                     "Worker task index, should be >= 0. task_index=0 is "
                     "the chief worker task the performs the variable "
                     "initialization")
flags.DEFINE_string("ps_hosts", ps_hosts,
                    "Comma-separated list of hostname:port pairs")
flags.DEFINE_string("worker_hosts", worker_hosts,
                    "Comma-separated list of hostname:port pairs")

# Training related flags
# Training related flags
flags.DEFINE_string("data_dir",
                    get_data_path(
                        dataset_name="malo/mnist",  # all mounted repo
                        local_root=ROOT_PATH_TO_LOCAL_DATA,
                        local_repo="mnist",
                        path=''
                    ),
                    "Path to store logs and checkpoints. It is recommended"
                    "to use get_logs_path() to define your logs directory."
                    "so that you can switch from local to clusterone without"
                    "changing your code."
                    "If you set your logs directory manually make sure"
                    "to use /logs/ when running on ClusterOne cloud.")
flags.DEFINE_string("log_dir",
                    get_logs_path(root=PATH_TO_LOCAL_LOGS),
                    "Path to dataset. It is recommended to use get_data_path()"
                    "to define your data directory.so that you can switch "
                    "from local to ClusterOne without changing your code."
                    "If you set the data directory manually makue sure to use"
                    "/data/ as root path when running on ClusterOne cloud.")

flags.DEFINE_integer("hidden1", 128,
                     "Number of units in the 1st hidden layer of the NN")
flags.DEFINE_integer("hidden2", 128,
                     "Number of units in the 2nd hidden layer of the NN")
flags.DEFINE_integer("batch_size", 100, "Training batch size")
flags.DEFINE_float("learning_rate", 0.01, "Learning rate")

FLAGS = flags.FLAGS

print(FLAGS.data_dir)
print(FLAGS.log_dir)


def device_and_target():
    # If FLAGS.job_name is not set, we're running single-machine TensorFlow.
    # Don't set a device.
    if FLAGS.job_name is None:
        print("Running single-machine training")
        return (None, "")

    # Otherwise we're running distributed TensorFlow.
    print("%s.%d  -- Running distributed training" % (FLAGS.job_name, FLAGS.task_index))
    if FLAGS.task_index is None or FLAGS.task_index == "":
        raise ValueError("Must specify an explicit `task_index`")
    if FLAGS.ps_hosts is None or FLAGS.ps_hosts == "":
        raise ValueError("Must specify an explicit `ps_hosts`")
    if FLAGS.worker_hosts is None or FLAGS.worker_hosts == "":
        raise ValueError("Must specify an explicit `worker_hosts`")

    cluster_spec = tf.train.ClusterSpec({
        "ps": FLAGS.ps_hosts.split(","),
        "worker": FLAGS.worker_hosts.split(","),
    })
    server = tf.train.Server(
        cluster_spec, job_name=FLAGS.job_name, task_index=FLAGS.task_index)
    if FLAGS.job_name == "ps":
        server.join()

    worker_device = "/job:worker/task:{}".format(FLAGS.task_index)
    # The device setter will automatically place Variables ops on separate
    # parameter servers (ps). The non-Variable ops will be placed on the workers.
    return (
        tf.train.replica_device_setter(
            worker_device=worker_device,
            cluster=cluster_spec),
        server.target,
    )


def main(unused_argv):
    if FLAGS.log_dir is None or FLAGS.log_dir == "":
        raise ValueError("Must specify an explicit `log_dir`")
    if FLAGS.data_dir is None or FLAGS.data_dir == "":
        raise ValueError("Must specify an explicit `data_dir`")

    print('Printing Flags')
    print(str(FLAGS))

    print(FLAGS.__flags)
    device, target = device_and_target()
    with tf.device(device):
        images = tf.placeholder(tf.float32, [None, 784], name='image_input')
        labels = tf.placeholder(tf.float32, [None], name='label_input')
        data = read_data_sets(FLAGS.data_dir,
                              one_hot=False,
                              fake_data=False)
        logits = mnist.inference(images, FLAGS.hidden1, FLAGS.hidden2)
        loss = mnist.loss(logits, labels)
        loss = tf.Print(loss, [loss], message="Loss = ")
        train_op = mnist.training(loss, FLAGS.learning_rate)

    with tf.train.MonitoredTrainingSession(
            master=target,
            is_chief=(FLAGS.task_index == 0),
            checkpoint_dir=FLAGS.log_dir) as sess:
        while not sess.should_stop():
            xs, ys = data.train.next_batch(FLAGS.batch_size, fake_data=False)
            sess.run(train_op, feed_dict={images: xs, labels: ys})


if __name__ == "__main__":
    tf.app.run()