有没有办法在TensorFlow中将堆叠的RNN并行化到多个GPU上？

Question

是否可以获取tf.scan操作的输出并将其直接流式传输到不同的GPU,从而有效地在两个GPU上并行运行两个堆叠的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 在输出处有一个“边界”，所有迭代都必须完成，然后输出张量才能被下一个操作读取。但是，您可以在不同的 GPU 上运行不同级别的 lstm 堆栈，并在扫描中获得帧并行性。编写您自己的 MultiRNNCell 版本，为每个 lstm 层使用单独的设备。

另外，您可能想使用 tf.nn.dynamic_rnn 而不是扫描。