Tensorflow - replicating model over multiple GPUs and sharing variables?

user_997338 0
Viewed 1086

I am trying to run a simple feed-forward network on multiple GPUs (to asynchronously updated shared weights).

However, I am unable to get the weights to be shared.

From the research I've done I just need to set reuse=True on variable_scope but that does not seem to be working:

for i_, gpu_id in enumerate(gpus):

    with tf.device(gpu_id):
        # [Build graph in here.]

        with variable_scope.variable_scope(variable_scope.get_variable_scope(), reuse=i_>0):

            x = tf.placeholder(tf.float32, [None, 784])
            W = tf.Variable(tf.zeros([784, 10]))
            b = tf.Variable(tf.zeros([10]))
            y = tf.nn.softmax(tf.matmul(x, W) + b)
            y_ = tf.placeholder(tf.float32, [None, 10])

            # More code..., see pastebin link below


# Start an interactive tensorflow session
sess = tf.Session()

# Initialize all variables associated with this session
sess.run(tf.initialize_all_variables())

The above is a sample of the code, in the full code (https://pastebin.com/i4NBnHHC) I show how training on one GPU does not update the weights on the other GPUs.

tensorflow distributed-computing parallel-processing gpu python
1 Answers

The simplest solution for you is to use the in-graph replication:

In-graph replication. In this approach, the client builds a single tf.Graph that contains one set of parameters (in tf.Variable nodes pinned to /job:ps); and multiple copies of the compute-intensive part of the model, each pinned to a different task in /job:worker.

For that you just place the parameters (placeholder and variables) on the CPU device:

# in-graph replication
import tensorflow as tf

num_gpus = 2

# place the initial data on the cpu
with tf.device('/cpu:0'):
    input_data = tf.Variable([[1., 2., 3.],
                              [4., 5., 6.],
                              [7., 8., 9.],
                              [10., 11., 12.]])
    b = tf.Variable([[1.], [1.], [2.]])

# split the data into chunks for each gpu
inputs = tf.split(input_data, num_gpus)
outputs = []

# loop over available gpus and pass input data
for i in range(num_gpus):
    with tf.device('/gpu:'+str(i)):
        outputs.append(tf.matmul(inputs[i], b))

# merge the results of the devices
with tf.device('/cpu:0'):
    output = tf.concat(outputs, axis=0)

# create a session and run
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print sess.run(output)

A more complex approach called between-graph replication is more preferred in tensorflow community, but would require a more sophisticated configuration with tf.train.ClusterSpec. You can see the example in their tutorial on distributed tensorflow.

Recommend this post that does comparison of different distribution settings.

user_7129950
Related