title: How to train mnist data with tensorflow ParameterServerStrategy distributed training?
body: <p>I'm trying to train the mnist dataset using the ParameterServerStrategy. As a beginner, I find the documentations to be confusing specially when it comes to the section &quot;Clusters in the real world&quot;. This is the docs that I'm following:<a href="https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/distribute/parameter_server_training.ipynb#scrollTo=zyby6M2Jqg6J&amp;uniqifier=1" rel="nofollow noreferrer">https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/distribute/parameter_server_training.ipynb#scrollTo=zyby6M2Jqg6J&amp;uniqifier=1</a> So far I have this:</p>
  <pre class="lang-py prettyprint-override"><code>
  #this is mnist_setup.py
  
  import os
  import tensorflow as tf
  import numpy as np
  
  def mnist_dataset(batch_size):
    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
    # The `x` arrays are in uint8 and have values in the [0, 255] range.
    # You need to convert them to float32 with values in the [0, 1] range.
    x_train = x_train / np.float32(255)
    y_train = y_train.astype(np.int64)
    train_dataset = tf.data.Dataset.from_tensor_slices(
        (x_train, y_train)).shuffle(60000).repeat().batch(batch_size)
    return train_dataset
  
  def build_and_compile_model():
    model = tf.keras.Sequential([
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(10, activation=tf.nn.softmax)
    ])
    model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        optimizer=tf.keras.optimizers.SGD(),
        metrics=['accuracy'])
    return model
  </code></pre>
  <pre class="lang-py prettyprint-override"><code>#this is main.py
  
  import os
  import json
  
  import tensorflow as tf
  import mnist_setup
  
  per_worker_batch_size = 64
  
  os.environ[&quot;TF_CONFIG&quot;] = json.dumps({
      &quot;cluster&quot;: {
          &quot;worker&quot;: ['ip_of_deeplearning_VM:port'], #worker1
          &quot;ps&quot;: ['ip_of_deeplearning_VM:port'], #worker2
          &quot;chief&quot;: ['ip_of_deeplearning_VM:port'] #masterz
      },
      &quot;task&quot;: {&quot;type&quot;: &quot;chief&quot;, &quot;index&quot;: 0}
  })
  
  cluster_spec = tf.train.ClusterSpec({
      'ps':['ip_of_deeplearning_VM:port'], #worker2
      'worker': ['ip_of_deeplearning_VM:port'] #worker1
  })
  
  cluster_resolver = tf.distribute.cluster_resolver.SimpleClusterResolver(cluster_spec, task_type=&quot;ps&quot;,task_id=0)
  
  tf_config = json.loads(os.environ['TF_CONFIG'])
  num_workers = len(tf_config['cluster']['worker'])
  
  strategy = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver)
  
  global_batch_size = per_worker_batch_size * num_workers
  multi_worker_dataset = mnist_setup.mnist_dataset(global_batch_size)
  
  with strategy.scope():
    # Model building/compiling need to be within `strategy.scope()`.
      multi_worker_model = mnist_setup.build_and_compile_model()
      
  print(&quot;chief gets called!&quot;)
  result = multi_worker_model.fit(multi_worker_dataset, epochs=3)
  </code></pre>
  <p>I copy these files to the worker and ps VM, change the index and run &quot;main.py&quot; on all of them at the same time. I get the message saying that the server was started at ip_address but that's about it. Would anyone please show me what I need to do in order to get this working?</p>