Есть ли способ распараллеливать сложенные RNN на несколько графических процессоров в TensorFlow?

Question

Возможно ли выполнить вывод операции tf.scan и передать ее прямо на другой графический процессор, эффективно выполняя два сложенных RNN на двух графических процессорах параллельно? Что-то вроде этого:

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)) 

Будет TensorFlow автоматически позаботится о том, оптимизации, или не будет его блокировать поток графика, пока список ys1 будет завершен.

Answer 1

К сожалению, tf.scan имеет «границу» на выходе, и все итерации должны быть выполнены до того, как выходные тензоры могут быть прочитаны следующими операциями. Тем не менее, вы можете запускать разные уровни стека lstm на разных графических процессорах и получать параллелизм кадров в ходе сканирования. Создайте собственную версию MultiRNNCell для использования отдельных устройств для каждого слоя lstm.

Также вы, вероятно, захотите использовать tf.nn.dynamic_rnn вместо сканирования.