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而不是扫描。

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

1 个答案: