Geteilter RNN-Speicherverbrauch gleichmäßig zwischen GPUs in TensorFlow

Question

Ich versuche, die meiste strategische Weise herauszufinden, die Speicherlast eines seq2seq-Netzwerks zwischen zwei GPUs gleichmäßig zu teilen.

Mit Faltungsnetzwerken ist die Aufgabe viel einfacher. Ich versuche jedoch herauszufinden, wie man den Speicherverbrauch von 2 Titan X maximieren kann. Das Ziel ist es, das größte Netzwerk zu erstellen, das die kombinierten 24 GB Speicher ermöglichen.

Eine Idee war es, jede RNN-Schicht in einer separaten GPU zu platzieren.

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

Die Backprop-Berechnungen erfordern jedoch eine erhebliche Menge an Speicher. Eine andere Idee besteht daher darin, den gesamten Vorwärtsdurchlauf auf einer GPU und den Rückwärtsdurchlauf auf der separaten GPU durchzuführen.

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

(jedoch GPU2 nimmt nach wie vor den größten Teil der Speicherauslastung)

Gibt es eine Möglichkeit, wie viel von der GPU-Speicher zu messen verwendet wird? Dies würde es uns ermöglichen, herauszufinden, wie man jede GPU maximieren kann, bevor sie "voll ist".

Sobald 2 GPUs verwendet werden, möchte ich schließlich vier verwenden. Ich denke jedoch, dass die Maximierung von 2 GPUs der erste Schritt ist.

Answer 1

Einstellung "colocate_gradients_with_ops" als TRUE möglicherweise funktionieren. Es ermöglicht GPU-Speicher gleichmäßig zugewiesen.

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)