如何将TensorFlow图(模型)分割成多个GPU以避免OOM？

Question

所以我有这个非常大和深的模型，我用TensorFlow r1.2实现，运行在一个12 GB内存的NVIDIA k40上。该模型由多个RNN、一组权值矩阵和嵌入矩阵以及偏置向量组成。当我启动培训项目时，首先要花2-3个小时来建模，然后由于OOM问题而崩溃。我试图将批处理大小减少到每批甚至一个数据样本，但仍然遇到了相同的问题。

如果我谷歌tensorflow muitlple gpu，我发现的例子主要是通过并行模型设计使用多个GPU，这意味着让每个GPU运行相同的图，并让CPU计算总梯度从而传播到每个参数。

我知道一个可能的解决方案是在内存更大的GPU上运行模型。但我不知道是否有一种方法可以将将的图形(模型)拆分为不同的部分--顺序为，并将它们分配给不同的GPU？

Answer 1

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

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