在小批量的多GPU上进行培训

时间:2018-08-08 14:35:13

标签: tensorflow

我在具有两个GPU(每个具有3 GB内存)的机器上运行TensorFlow。我的批量大小只有2GB,因此可以容纳一个GPU。使用两个GPU(使用CUDA_VISIBLE_DEVICES)进行训练有什么意义吗?如果我这样做了,TensorFlow将如何分配培训?

1 个答案:

答案 0 :(得分:0)

关于内存:我假设您的意思是一个数据批为2GB。但是,Tensorflow还需要内存来存储变量以及隐藏层结果等(以计算梯度)。因此,内存是否足够还取决于您的特定型号。最好的选择是仅尝试使用一个GPU,看看程序是否由于内存错误而崩溃。

关于分发:Tensorflow完全不会自动执行此操作。每个操作都放置在某些设备上。默认情况下,如果有任意数量的GPU,则所有与GPU兼容的操作将放置在第一个GPU上,其余放置在CPU上。尽管Tensorflow默认情况下会保留所有GPU上的所有内存。

您应该看看the GPU guide on the Tensorflow website。最重要的是,您可以使用with tf.device上下文管理器将操作放置在其他GPU上。使用此方法,可以将批处理分成X个块(X = GPU数量),并在每台设备上定义模型,每次将各自的块作为输入并确保重用变量。

如果您使用的是tf.Estimator,则有些信息in this question。仅使用两个简单的包装程序就可以很容易地进行分布式执行,但是我个人却无法成功使用它(相当慢,并且会因段错误随机崩溃)。