如何将TensorFlow图（模型）拆分到多个GPU上以避免OOM？-Java 学习之路

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

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

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

1 回答

1
using GPUs上的官方指南向您展示了"Using multiple GPUs"中的示例 . 您只需要在不同的tf.device上下文中创建操作;节点仍将添加到同一图表中，但它们将使用指示应在何处运行的设备指令进行注释 . 例如：
```
with tf.device("/gpu:0"):
    net0 = make_subnet0()
with tf.device("/gpu:1"):
    net1 = make_subnet1()
result = combine_subnets(net0, net1)
```
回复于 2024-04-29T00:42:03+08:00

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

1 回答

相关问题