在 TensorFlow 中的 GPU 之间平均分配 RNN 内存消耗

Posted 2023-02-16

【中文标题】在 TensorFlow 中的 GPU 之间平均分配 RNN 内存消耗

【英文标题】：Split RNN Memory Consumption Evenly Between GPUs in TensorFlow

【发布时间】：2017-02-07 23:26:37

【问题描述】：

我正在尝试找出最具战略意义的方法，以在两个 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。

一旦使用了 2 个 GPU，我最终会想要使用 4 个。但是，我认为最大化 2 个 GPU 是第一步。

【参考方案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)