Split RNN Mức tiêu thụ bộ nhớ Đồng đều giữa các GPU trong TensorFlow

Question

Tôi đang cố gắng tìm ra cách chiến lược nhất để phân chia đồng đều tải bộ nhớ của mạng seq2seq giữa hai GPU.

Với mạng chuyển đổi, nhiệm vụ sẽ dễ dàng hơn nhiều. Tuy nhiên, tôi đang cố gắng tìm ra cách tối đa hóa việc sử dụng bộ nhớ của 2 Titan X. Mục tiêu là xây dựng mạng lớn nhất mà bộ nhớ 24GB kết hợp sẽ cho phép.

Một ý tưởng là đặt từng lớp RNN vào một GPU riêng biệt.

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

Tuy nhiên, tính toán backprop yêu cầu một lượng bộ nhớ đáng kể. Do đó, một ý tưởng khác là thực hiện toàn bộ chuyển tiếp trên một GPU và đường truyền ngược trên GPU riêng biệt.

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

(Tuy nhiên, GPU2 vẫn mất hầu hết các tải bộ nhớ)

Có cách nào để đo lường bao nhiêu bộ nhớ GPU đang được sử dụng? Điều này sẽ cho phép chúng tôi tìm ra cách tối đa hóa từng GPU trước khi nó được "lấp đầy".

Sau khi sử dụng 2 GPU, cuối cùng tôi sẽ muốn sử dụng bốn GPU. Tuy nhiên, tôi nghĩ tối đa 2 GPU là bước đầu tiên.

Answer 1

Đặt "colocate_gradients_with_ops" là TRUE có thể hoạt động. Nó cho phép bộ nhớ GPU được phân bổ đồng đều.

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)