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를 최대화하는 방법을 파악할 수 있습니다.

GPU 2 개가 사용되면 결국 4 개를 사용하려고합니다. 그러나 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)