Is there a way to parallelize stacked RNNs over multiple GPUs in TensorFlow?

Question

Is it possible to take the output of a tf.scan operation and stream it directly to a different GPU, effectively running two stacked RNNs on two GPUs in parallel? Something like this:

cell1 = tf.nn.rnn_cell.MultiRNNCell(..)
cell2 = tf.nn.rnn_cell.MultiRNNCell(..)

with tf.device("/gpu:0"):
  ys1 = tf.scan(lambda a, x: cell1(x, a[1]), inputs,
          initializer=(tf.zeros([batch_size, state_size]), init_state))

with tf.device("/gpu:1"):
  ys2 = tf.scan(lambda a, x: cell2(x, a[1]), ys1,
          initializer=(tf.zeros([batch_size, state_size]), init_state))

Will TensorFlow automatically take care of that optimization, or will it block the graph flow until the list ys1 is finalized.

Answer 1

Unfortunately, tf.scan has a "boundary" at the output, and all iterations have to complete before the output tensor can be read by the next operations. However, you can run the different levels of your lstm stack on different GPUs, and get frame parallelism within a scan. Write your own version of MultiRNNCell to use separate devices for each lstm layer.

Also you probably want to use tf.nn.dynamic_rnn instead of scan.