title: MonitoredTrainingSession save and restore model
body: <p>I'm trying to extend the example <a href="https://www.tensorflow.org/deploy/distributed" rel="nofollow noreferrer">https://www.tensorflow.org/deploy/distributed</a> outlined here but I'm having trouble saving the model. I'm running this in docker container available at <code>gcr.io/tensorflow/tensorflow:1.5.0-gpu-py3</code>. I started two processes one for 'ps' and one for 'worker' and the ps process is simply this code:</p>
  
  <pre><code>import tensorflow as tf
  def main(_):
     cluster = tf.train.ClusterSpec({"ps":["localhost:2222"],"worker":["localhost:2223"]})
     server = tf.train.Server(cluster, job_name="ps", task_index=0)
     server.join()
  if __name__ == "__main__":
     tf.app.run()
  </code></pre>
  
  <p>The worker code is the following and is based on the mnist examples and the distributed article above:</p>
  
  <pre><code>import tensorflow as tf
  from tensorflow.examples.tutorials.mnist import input_data
  
  data_dir = "/data"
  checkpoint_dir = "/tmp/train_logs"
  
  def main(_):
     cluster = tf.train.ClusterSpec({"ps":["localhost:2222"],"worker":["localhost:2223"]})
     server = tf.train.Server(cluster, job_name="worker", task_index=0)
     mnist = input_data.read_data_sets(data_dir, one_hot=True)
  
     with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:0", cluster=cluster)):
       x = tf.placeholder(tf.float32, [None,784], name="x_input")
       W = tf.Variable(tf.zeros([784,10]))
       b = tf.Variable(tf.zeros([10]))
       y = tf.placeholder(tf.float32, [None,10])
       model = tf.matmul(x, W) + b
       cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=y, logits=model))
       global_step = tf.train.get_or_create_global_step()
       train_op = tf.train.GradientDescentOptimizer(0.5).minimize(cost, global_step=global_step)
       prediction = tf.equal(tf.argmax(model,1), tf.argmax(y,1))
       accuracy = tf.reduce_mean(tf.cast(prediction, tf.float32))
  
    hooks = [tf.train.StopAtStepHook(last_step=101)]
    with tf.train.MonitoredTrainingSession(master=server.target, is_chief=True, checkpoint_dir=checkpoint_dir, hooks=hooks) as sess:
       while not sess.should_stop():
          batch_xs, batch_ys = mnist.train.next_batch(1000)
          sess.run(train_op, feed_dict={x: batch_xs, y: batch_ys})
  
    latest_checkpoint = tf.train.latest_checkpoint(checkpoint_dir)
    #saver = tf.train.Saver()
    saver = tf.train.import_meta_graph(latest_checkpoint+".meta", clear_devices=True)
    with tf.Session() as sess:
       saver.restore(sess,latest_checkpoint) # "/tmp/train_logs/model.ckpt"
       acc = sess.run(accuracy, feed_dict={x: mnist.test.images,y: mnist.test.labels});
       print("Test accuracy = "+"{:5f}".format(acc))
  
  if __name__ == "__main__":
     tf.app.run()
  </code></pre>
  
  <p>The examples I've found all seem to end without showing how to use the model. The above code fails on the <code>saver.restore()</code> line with the following error:</p>
  
  <pre><code>InvalidArgumentError (see above for traceback): Cannot assign a device for operation 'save/RestoreV2_2': 
  Operation was explicitly assigned to /job:ps/task:0/device:CPU:0 
  but available devices are [ /job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:GPU:0 ].
  Make sure the device specification refers to a valid device.
  </code></pre>
  
  <p>Also, as shown above I tried both <code>saver = tf.train.Saver()</code> and <code>saver = tf.train.import_meta_graph(latest_checkpoint+".meta", clear_devices=True)</code> with no success. Same error is shown in either case.</p>
  
  <p>I don't really understand the <code>with tf.device(...):</code> statement. In one iteration I commented out this line (and unindented the statements below it) and the code ran without errors. But I think this is not correct and would like to understand the correct way for this to work.</p>