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

Question

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

对于卷积网络来说，任务要简单得多。但是，我试图找出如何最大限度地利用2 Titan X的内存。我们的目标是建立最大的网络，合并的24GB内存将允许。

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

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

但是，反向计算需要大量的内存。因此，另一个想法是在一个GPU上执行整个正向传递，并在单独的GPU上执行反向传递。

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

（然而，GPU2仍占大多数内存负载）

有什么办法来衡量GPU如何内存中有多少被使用？这将使我们能够弄清楚如何在每个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)