Có cách nào để song song RNN xếp chồng lên nhau trên nhiều GPU trong TensorFlow không?

Question

Có thể lấy đầu ra của hoạt động tf.scan và truyền trực tiếp tới một GPU khác, chạy hai RNN được chồng lên nhau trên hai GPU song song hay không? Một cái gì đó như thế này:

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

sẽ TensorFlow tự động chăm sóc tối ưu hóa đó, hoặc nó sẽ chặn dòng chảy đồ thị cho đến khi danh sách ys1 được hoàn thành.

Answer 1

Thật không may, tf.scan có một "ranh giới" ở đầu ra, và tất cả các lần lặp lại phải hoàn thành trước khi các thiết bị tiếp theo có thể đọc được thiết bị đầu ra. Tuy nhiên, bạn có thể chạy các cấp độ khác nhau của chồng lstm trên các GPU khác nhau và nhận được sự tương đồng khung trong một lần quét. Viết phiên bản MultiRNNCell của riêng bạn để sử dụng các thiết bị riêng biệt cho mỗi lớp lstm.

Ngoài ra, bạn có thể muốn sử dụng tf.nn.dynamic_rnn thay vì quét.