title: How to use tf.keras Sequential with tf.distribute.ParameterServerStrategy and tf.train.MonitoredSession?
body: <p>I'm trying to set up a really easy Minimal Working Example for the following: Use a model built with tf.keras in a tf.train.MonitoredSession using a tf.Server with a tf.distribute.ParameterServerStrategy.</p>
  
  <p>My goal in the end is to use a tf.keras model in a distributed environment using two workers each having one GPU and a parameter server.</p>
  
  <p>The model is built according to the example and documentation found here: <a href="https://www.tensorflow.org/versions/r1.12/api_docs/python/tf/keras/models/Sequential" rel="nofollow noreferrer">https://www.tensorflow.org/versions/r1.12/api_docs/python/tf/keras/models/Sequential</a></p>
  
  <p>The parameter server strategy is used according to the documentation found here: <a href="https://www.tensorflow.org/versions/r1.12/api_docs/python/tf/contrib/distribute/ParameterServerStrategy" rel="nofollow noreferrer">https://www.tensorflow.org/versions/r1.12/api_docs/python/tf/contrib/distribute/ParameterServerStrategy</a></p>
  
  <p>The overall setup including the device placement and the use of a MonitoredSession is taken from: <a href="https://github.com/tensorflow/examples/blob/master/community/en/docs/deploy/distributed.md" rel="nofollow noreferrer">https://github.com/tensorflow/examples/blob/master/community/en/docs/deploy/distributed.md</a></p>
  
  <p>I'm already using the "allow_soft_placement" option and I'm "emulating" a distributed setup on my local machine having only a single CPU, since there are different problems in the real distributed setup which I'm trying to solve by using a MonitoredSession where the variable initialization is handled automatically.</p>
  
  <p>This Code works with a "normal" (not monitored) tf.Session and variable initialization - global, local, model variables and tables etc.</p>
  
  <p>The line which unfreezes the graph is necessary to be able to use a tf.data.Dataset in the tf.keras.Model's fit function, since an iterator has to be created - which causes an error in a frozen graph.</p>
  
  <p>This the code I'm trying to run. I use tensorflow 1.12.0 and python 3.6.7. I've also tried python 2.7, same result.</p>
  
  <p>The code requires no setup besides installing tensorflow.</p>
  
  <pre class="lang-py prettyprint-override"><code>import sys
  import tensorflow as tf
  
  def main(argv):
  
    # Create local cluster config for run_local_server.sh script.
    cluster = tf.train.ClusterSpec({"worker": ["localhost:2222"], "ps": ["localhost:2223"]})
    task = 0
    job = str(argv[0])
  
    # Number of GPUs per worker
    GPU_PER_WORKER = 0
  
    config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)
    server = tf.train.Server(cluster, job_name=job, task_index=task,config=config)
  
    strategy = tf.contrib.distribute.ParameterServerStrategy(num_gpus_per_worker=GPU_PER_WORKER)
    strategy.configure(session_config=config, cluster_spec=cluster,task_type=job,task_id=task)
  
    with strategy.scope():
      (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
  
      x_train = x_train.astype('float32') / 255
      x_test = x_test.astype('float32') / 255
  
      # Reshape input data from (28, 28) to (28, 28, 1)
      w, h = 28, 28
      x_train = x_train.reshape(x_train.shape[0], w, h, 1)
      x_test = x_test.reshape(x_test.shape[0], w, h, 1)
  
      # One-hot encode the labels
      y_train = tf.keras.utils.to_categorical(y_train, 10)
      y_test = tf.keras.utils.to_categorical(y_test, 10)
  
      train_ds = tf.data.Dataset.zip((tf.data.Dataset.from_tensor_slices(x_train),tf.data.Dataset.from_tensor_slices(y_train))).repeat().shuffle(60000).batch(10)
      val_ds = tf.data.Dataset.zip((tf.data.Dataset.from_tensor_slices(x_test),tf.data.Dataset.from_tensor_slices(y_test))).repeat().shuffle(10000).batch(10)
  
      with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % task,cluster=cluster)):
        model = tf.keras.models.Sequential()
  
        conv0 = tf.keras.layers.Conv2D(filters=16, data_format='channels_last', padding="valid", kernel_size=4, strides=1, input_shape=(28,28,1), activation=tf.keras.activations.relu)
        model.add(conv0)
  
        flatten = tf.keras.layers.Flatten()
        model.add(flatten)
  
        dense1 = tf.keras.layers.Dense(10, activation=tf.keras.activations.softmax)
        model.add(dense1)
  
        model.compile(tf.contrib.optimizer_v2.AdamOptimizer(0.001), loss=tf.keras.metrics.mean_absolute_error,metrics=['accuracy'],distribute=strategy)   
  
      if job == "ps":
        server.join()
      elif job == "worker":
        with tf.train.MonitoredSession(session_creator=tf.train.ChiefSessionCreator(master=server.target,config=config)) as sess:        
          sess.graph._unsafe_unfinalize()
          history = model.fit(x=train_ds, validation_data=val_ds, validation_steps=1000, steps_per_epoch=100, epochs=60)
  
  if __name__ == "__main__":
    main(sys.argv[1:])
  </code></pre>
  
  <p>The code requires no extensive setup since the dataset is loaded from the web and converted to a tf.data.Dataset since this is how I want to organize my pipeline with the real data. The MNIST data setup example is taken from 
  <a href="https://www.kaggle.com/margaretmz/mnist-with-tf-keras" rel="nofollow noreferrer">https://www.kaggle.com/margaretmz/mnist-with-tf-keras</a>.</p>
  
  <p>I expect the code to not fail due to wrong variable or operation placement since I'm basically leaving all these decisions to the implementation of tensorflow by using <code>strategy.scope()</code> and <code>tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % task,cluster=cluster))</code></p>