1 # Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 # ============================================================================== 15 """Distributed MNIST training and validation, with model replicas. 16 17 A simple softmax model with one hidden layer is defined. The parameters 18 (weights and biases) are located on one parameter server (ps), while the ops 19 are executed on two worker nodes by default. The TF sessions also run on the 20 worker node. 21 Multiple invocations of this script can be done in parallel, with different 22 values for --task_index. There should be exactly one invocation with 23 --task_index, which will create a master session that carries out variable 24 initialization. The other, non-master, sessions will wait for the master 25 session to finish the initialization before proceeding to the training stage. 26 27 The coordination between the multiple worker invocations occurs due to 28 the definition of the parameters on the same ps devices. The parameter updates 29 from one worker is visible to all other workers. As such, the workers can 30 perform forward computation and gradient calculation in parallel, which 31 should lead to increased training speed for the simple model. 32 """ 33 34 from __future__ import absolute_import 35 from __future__ import division 36 from __future__ import print_function 37 38 import math 39 import sys 40 import tempfile 41 import time 42 43 import tensorflow as tf 44 from tensorflow.examples.tutorials.mnist import input_data 45 46 flags = tf.app.flags 47 flags.DEFINE_string("data_dir", "/tmp/mnist-data", 48 "Directory for storing mnist data") 49 flags.DEFINE_boolean("download_only", False, 50 "Only perform downloading of data; Do not proceed to " 51 "session preparation, model definition or training") 52 flags.DEFINE_integer("task_index", None, 53 "Worker task index, should be >= 0. task_index=0 is " 54 "the master worker task the performs the variable " 55 "initialization ") 56 flags.DEFINE_integer("num_gpus", 1, "Total number of gpus for each machine." 57 "If you don't use GPU, please set it to '0'") 58 flags.DEFINE_integer("replicas_to_aggregate", None, 59 "Number of replicas to aggregate before parameter update" 60 "is applied (For sync_replicas mode only; default: " 61 "num_workers)") 62 flags.DEFINE_integer("hidden_units", 100, 63 "Number of units in the hidden layer of the NN") 64 flags.DEFINE_integer("train_steps", 200, 65 "Number of (global) training steps to perform") 66 flags.DEFINE_integer("batch_size", 100, "Training batch size") 67 flags.DEFINE_float("learning_rate", 0.01, "Learning rate") 68 flags.DEFINE_boolean( 69 "sync_replicas", False, 70 "Use the sync_replicas (synchronized replicas) mode, " 71 "wherein the parameter updates from workers are aggregated " 72 "before applied to avoid stale gradients") 73 flags.DEFINE_boolean( 74 "existing_servers", False, "Whether servers already exists. If True, " 75 "will use the worker hosts via their GRPC URLs (one client process " 76 "per worker host). Otherwise, will create an in-process TensorFlow " 77 "server.") 78 flags.DEFINE_string("ps_hosts", "localhost:2222", 79 "Comma-separated list of hostname:port pairs") 80 flags.DEFINE_string("worker_hosts", "localhost:2223,localhost:2224", 81 "Comma-separated list of hostname:port pairs") 82 flags.DEFINE_string("job_name", None, "job name: worker or ps") 83 84 FLAGS = flags.FLAGS 85 86 IMAGE_PIXELS = 28 87 88 89 def main(unused_argv): 90 mnist = input_data.read_data_sets(FLAGS.data_dir, one_hot=True) 91 if FLAGS.download_only: 92 sys.exit(0) 93 94 if FLAGS.job_name is None or FLAGS.job_name == "": 95 raise ValueError("Must specify an explicit `job_name`") 96 if FLAGS.task_index is None or FLAGS.task_index == "": 97 raise ValueError("Must specify an explicit `task_index`") 98 99 print("job name = %s" % FLAGS.job_name) 100 print("task index = %d" % FLAGS.task_index) 101 102 #Construct the cluster and start the server 103 ps_spec = FLAGS.ps_hosts.split(",") 104 worker_spec = FLAGS.worker_hosts.split(",") 105 106 # Get the number of workers. 107 num_workers = len(worker_spec) 108 109 cluster = tf.train.ClusterSpec({"ps": ps_spec, "worker": worker_spec}) 110 111 if not FLAGS.existing_servers: 112 # Not using existing servers. Create an in-process server. 113 server = tf.train.Server( 114 cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index) 115 if FLAGS.job_name == "ps": 116 server.join() 117 118 is_chief = (FLAGS.task_index == 0) 119 if FLAGS.num_gpus > 0: 120 # Avoid gpu allocation conflict: now allocate task_num -> #gpu 121 # for each worker in the corresponding machine 122 gpu = (FLAGS.task_index % FLAGS.num_gpus) 123 worker_device = "/job:worker/task:%d/gpu:%d" % (FLAGS.task_index, gpu) 124 elif FLAGS.num_gpus == 0: 125 # Just allocate the CPU to worker server 126 cpu = 0 127 worker_device = "/job:worker/task:%d/cpu:%d" % (FLAGS.task_index, cpu) 128 # The device setter will automatically place Variables ops on separate 129 # parameter servers (ps). The non-Variable ops will be placed on the workers. 130 # The ps use CPU and workers use corresponding GPU 131 with tf.device( 132 tf.train.replica_device_setter( 133 worker_device=worker_device, 134 ps_device="/job:ps/cpu:0", 135 cluster=cluster)): 136 global_step = tf.Variable(0, name="global_step", trainable=False) 137 138 # Variables of the hidden layer 139 hid_w = tf.Variable( 140 tf.truncated_normal( 141 [IMAGE_PIXELS * IMAGE_PIXELS, FLAGS.hidden_units], 142 stddev=1.0 / IMAGE_PIXELS), 143 name="hid_w") 144 hid_b = tf.Variable(tf.zeros([FLAGS.hidden_units]), name="hid_b") 145 146 # Variables of the softmax layer 147 sm_w = tf.Variable( 148 tf.truncated_normal( 149 [FLAGS.hidden_units, 10], 150 stddev=1.0 / math.sqrt(FLAGS.hidden_units)), 151 name="sm_w") 152 sm_b = tf.Variable(tf.zeros([10]), name="sm_b") 153 154 # Ops: located on the worker specified with FLAGS.task_index 155 x = tf.placeholder(tf.float32, [None, IMAGE_PIXELS * IMAGE_PIXELS]) 156 y_ = tf.placeholder(tf.float32, [None, 10]) 157 158 hid_lin = tf.nn.xw_plus_b(x, hid_w, hid_b) 159 hid = tf.nn.relu(hid_lin) 160 161 y = tf.nn.softmax(tf.nn.xw_plus_b(hid, sm_w, sm_b)) 162 cross_entropy = -tf.reduce_sum(y_ * tf.log(tf.clip_by_value(y, 1e-10, 1.0))) 163 164 opt = tf.train.AdamOptimizer(FLAGS.learning_rate) 165 166 if FLAGS.sync_replicas: 167 if FLAGS.replicas_to_aggregate is None: 168 replicas_to_aggregate = num_workers 169 else: 170 replicas_to_aggregate = FLAGS.replicas_to_aggregate 171 172 opt = tf.train.SyncReplicasOptimizer( 173 opt, 174 replicas_to_aggregate=replicas_to_aggregate, 175 total_num_replicas=num_workers, 176 name="mnist_sync_replicas") 177 178 train_step = opt.minimize(cross_entropy, global_step=global_step) 179 180 if FLAGS.sync_replicas: 181 local_init_op = opt.local_step_init_op 182 if is_chief: 183 local_init_op = opt.chief_init_op 184 185 ready_for_local_init_op = opt.ready_for_local_init_op 186 187 # Initial token and chief queue runners required by the sync_replicas mode 188 chief_queue_runner = opt.get_chief_queue_runner() 189 sync_init_op = opt.get_init_tokens_op() 190 191 init_op = tf.global_variables_initializer() 192 train_dir = tempfile.mkdtemp() 193 194 if FLAGS.sync_replicas: 195 sv = tf.train.Supervisor( 196 is_chief=is_chief, 197 logdir=train_dir, 198 init_op=init_op, 199 local_init_op=local_init_op, 200 ready_for_local_init_op=ready_for_local_init_op, 201 recovery_wait_secs=1, 202 global_step=global_step) 203 else: 204 sv = tf.train.Supervisor( 205 is_chief=is_chief, 206 logdir=train_dir, 207 init_op=init_op, 208 recovery_wait_secs=1, 209 global_step=global_step) 210 211 sess_config = tf.ConfigProto( 212 allow_soft_placement=True, 213 log_device_placement=False, 214 device_filters=["/job:ps", 215 "/job:worker/task:%d" % FLAGS.task_index]) 216 217 # The chief worker (task_index==0) session will prepare the session, 218 # while the remaining workers will wait for the preparation to complete. 219 if is_chief: 220 print("Worker %d: Initializing session..." % FLAGS.task_index) 221 else: 222 print("Worker %d: Waiting for session to be initialized..." % 223 FLAGS.task_index) 224 225 if FLAGS.existing_servers: 226 server_grpc_url = "grpc://" + worker_spec[FLAGS.task_index] 227 print("Using existing server at: %s" % server_grpc_url) 228 229 sess = sv.prepare_or_wait_for_session(server_grpc_url, config=sess_config) 230 else: 231 sess = sv.prepare_or_wait_for_session(server.target, config=sess_config) 232 233 print("Worker %d: Session initialization complete." % FLAGS.task_index) 234 235 if FLAGS.sync_replicas and is_chief: 236 # Chief worker will start the chief queue runner and call the init op. 237 sess.run(sync_init_op) 238 sv.start_queue_runners(sess, [chief_queue_runner]) 239 240 # Perform training 241 time_begin = time.time() 242 print("Training begins @ %f" % time_begin) 243 244 local_step = 0 245 while True: 246 # Training feed 247 batch_xs, batch_ys = mnist.train.next_batch(FLAGS.batch_size) 248 train_feed = {x: batch_xs, y_: batch_ys} 249 250 _, step = sess.run([train_step, global_step], feed_dict=train_feed) 251 local_step += 1 252 253 now = time.time() 254 print("%f: Worker %d: training step %d done (global step: %d)" % 255 (now, FLAGS.task_index, local_step, step)) 256 257 if step >= FLAGS.train_steps: 258 break 259 260 time_end = time.time() 261 print("Training ends @ %f" % time_end) 262 training_time = time_end - time_begin 263 print("Training elapsed time: %f s" % training_time) 264 265 # Validation feed 266 val_feed = {x: mnist.validation.images, y_: mnist.validation.labels} 267 val_xent = sess.run(cross_entropy, feed_dict=val_feed) 268 print("After %d training step(s), validation cross entropy = %g" % 269 (FLAGS.train_steps, val_xent)) 270 271 272 if __name__ == "__main__": 273 tf.app.run() 274