Разделение памяти RNN равномерно между графическими процессорами в TensorFlow

Question

Я пытаюсь найти наиболее стратегический способ равномерного разделения нагрузки на память сети seq2seq между двумя графическими процессорами.

С сверточными сетями задача намного проще. Тем не менее, я пытаюсь выяснить, как максимально использовать память 2 Titan X. Цель состоит в том, чтобы построить самую большую сеть, которая позволит объединить 24 ГБ памяти.

Одна из идей заключалась в том, чтобы разместить каждый слой RNN в отдельном графическом процессоре.

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

Однако вычисления backprop требуют значительного объема памяти. Поэтому другая идея состоит в том, чтобы сделать весь передний проход на одном графическом процессоре и обратном проходе на отдельный графический процессор.

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

(Тем не менее, по-прежнему GPU2 занимает большая часть нагрузки на памяти)

Есть ли способ, чтобы определить, сколько памяти GPU используются? Это позволило бы нам выяснить, как максимизировать каждый графический процессор, прежде чем он будет «заполнен».

Как только 2 GPU используются, я в конце концов захочу использовать четыре. Тем не менее, я думаю, что максимизация 2 GPU - это первый шаг.

Answer 1

Настройка «colocate_gradients_with_ops» как ИСТИНА, возможно, работает. Он позволяет распределять память 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)