如何将 TensorFlow 图（模型）拆分到多个 GPU 上以避免 OOM？

Posted 2023-02-16

【中文标题】如何将 TensorFlow 图（模型）拆分到多个 GPU 上以避免 OOM？

【英文标题】：How to split TensorFlow graph (model) onto multiple GPUs to avoid OOM?

【发布时间】：2017-07-27 09:19:54

【问题描述】：

因此，我使用 TensorFlow r1.2 实现了这个非常大且深的模型，在具有 12 GB 内存的 NVIDIA Tesla k40 上运行。该模型由几个 RNN、一堆权重和嵌入矩阵以及偏置向量组成。当我启动训练程序时，首先构建模型需要大约 2-3 个小时，然后由于 OOM 问题而崩溃。我试图将批次大小减少到每批次 1 个数据样本，但仍然遇到了同样的问题。

如果我 google tensorflow muitlple gpu，我发现的示例主要集中在通过 并行 模型设计来利用多个 GPU，这意味着让每个 GPU 运行相同 绘制图形并让 CPU 计算总梯度，从而传播回每个参数。

我知道一种可能的解决方案是在具有更大内存的 GPU 上运行模型。但我想知道是否有办法将我的图（模型）按顺序拆分分成不同的部分并将它们分配给不同的 GPU？

【参考方案1】：

using GPUs 上的官方指南在“使用多个 GPU”中向您展示了该示例。您只需要在不同的tf.device 上下文中创建操作；节点仍将添加到同一图表中，但它们将使用设备指令进行注释，指示它们应在何处运行。例如：

with tf.device("/gpu:0"):
    net0 = make_subnet0()
with tf.device("/gpu:1"):
    net1 = make_subnet1()
result = combine_subnets(net0, net1)