Is it possible to split a network into several GPUs in a tensor stream?

Question

I plan to launch a very large recurrent network (for example, 2048x5), is it possible to define one layer on one GPU in a tensor stream? How should I implement the model for maximum efficiency. I understand that there is overhead for exchanging data between the GPU or GPU-CPU-GPU.

Answer 1

Dividing a large model into several GPUs is certainly possible in TensorFlow, but the optimal solution is a complex research problem. In general, you will need to do the following:

• Wrap large contiguous areas of your code in a block with tf.device(...): calling different GPUs:

   with tf.device("/gpu:0"): # Define first layer. with tf.device("/gpu:1"): # Define second layer. # Define other layers, etc. 

• When creating the optimizer, pass the optional colocate_gradients_with_ops=True argument to the optimizer.minimize() method:

   loss = ... optimizer = tf.train.AdaGradOptimizer(0.01) train_op = optimizer.minimize(loss, colocate_gradients_with_ops=True) 

• (optional.) Perhaps you need to enable "soft placement" in tf.ConfigProto when creating tf.Session , if any of the operations in your model cannot be performed on the GPU:

   config = tf.ConfigProto(allow_soft_placement=True) sess = tf.Session(config=config)