Is it possible to split a network across multiple GPUs in tensorflow?

Question

I plan to run a very large recurrent network (eg 2048x5), is it possible to define one layer at one GPU in tensorflow? How should I implement the model to achieve the best efficiency. I understand there is overhead for inter-GPU or GPU-CPU-GPU communication.

Answer 1

Splitting a large model across multiple GPUs is certainly possible in TensorFlow, but doing it optimally is a hard research problem. In general, you will need to do the following:

1. Wrap large contiguous regions of your code in a with tf.device(...): block, naming the different GPUs:

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

2. When building your optimizer, pass the optional argument colocate_gradients_with_ops=True to the optimizer.minimize() method:

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

3. (Optionally.) You may need to enable "soft placement" in the tf.ConfigProto when you create your tf.Session , if any of the operations in your model cannot run on GPU:

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