Skip to content

Instantly share code, notes, and snippets.

@benman1
Last active July 24, 2019 14:59
Show Gist options
  • Save benman1/318bcad0f1a9bba1046a7d8eb452603b to your computer and use it in GitHub Desktop.
Save benman1/318bcad0f1a9bba1046a7d8eb452603b to your computer and use it in GitHub Desktop.
import tensorflow as tf
# run this script from the master node
cluster = tf.train.ClusterSpec({"distributed": ["192.168.102.90:1111", "192.168.102.90:1111"]})
x = tf.constant(2)
with tf.device("/job:local/replica:0/task:0/device:GPU:0"): #"/job:distributed/task:1"):
y2 = x - 66
with tf.device("/job:local/replica:0/task:1/device:GPU:0"): #"/job:distributed/task:0"):
y1 = x + 300
y = y1 + y2
with tf.Session("grpc://localhost:1111") as sess:
result = sess.run(y)
print(result)
import sys
task_number = int(sys.argv[1])
import tensorflow as tf
# use this script to start the nodes
# for the first node:
# python start_node 0
cluster = tf.train.ClusterSpec({"local": ["192.168.102.20:1111", "192.168.102.90:1111"]})
server = tf.train.Server(cluster, job_name="local", task_index=task_number)
print("Starting server #{}".format(task_number))
server.start()
server.join()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment