在TensorFlow中,在GPU之间均匀分割RNN内存消耗

Question

我试图找出在两个GPU之间均匀分配seq2seq网络的内存负载的最具战略性的方法.

使用卷积网络,任务更容易.但是,我正在试图弄清楚如何最大化2 Titan X的内存使用量.目标是构建最大的24GB内存允许的网络.

一个想法是将每个RNN层放在单独的GPU中.

GPU1 --> RNN Layer 1 & Backward Pass
GPU2 --> RNN Layer 2,3,4

但是,backprop计算需要大量内存.因此,另一个想法是在一个GPU上进行整个前向传递,在单独的GPU上进行后向传递.

GPU1 --> Forward Pass
GPU2 --> Backward Pass

(但是,GPU2仍占用大部分内存负载)

有没有办法测量使用了多少GPU内存？这将使我们能够弄清楚如何在每个GPU"填满"之前最大化它们.

一旦使用了2个GPU,我最终会想要使用4个GPU.但是,我认为最大化2个GPU是第一步.

Answer 1

将“colocate_gradients_with_ops”设置为 TRUE 也许可行。它允许均匀分配 GPU 内存。

optimizer = tf.train.AdamOptimizer(learning_rate)
gvs = optimizer.compute_gradients(loss, colocate_gradients_with_ops=True)
train_op = optimizer.apply_gradients(gvs, global_step=self.global_step)