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