tensorflow - 在小批量的多 GPU 上进行训练

Question

我在一台有两个 GPU 的机器上运行 TensorFlow，每个 GPU 有 3 GB 内存。我的批量大小只有 2GB，因此可以放在一个 GPU 上。CUDA_VISIBLE_DEVICES使用两个 GPU（使用）进行训练有什么意义吗？如果我这样做了，TensorFlow 将如何分配训练？

标签： tensorflow

Answer 1

关于内存：我假设您的意思是一个数据批次是 2GB。然而，Tensorflow 也需要内存来存储变量以及隐藏层结果等（计算梯度）。因此，内存是否足够也取决于您的特定型号。最好的办法是尝试使用一个 GPU，看看程序是否由于内存错误而崩溃。

关于分发：Tensorflow 根本不会自动执行此操作。每个操作都放置在某个设备上。默认情况下，如果您有任意数量的 GPU 可用，所有与 GPU 兼容的操作都将放在第一个 GPU 上，其余的放在 CPU 上。尽管默认情况下 TensorFlow 会在所有 GPU 上保留所有内存。

您应该查看Tensorflow 网站上的 GPU 指南。最重要的是您可以使用with tf.device上下文管理器将操作放置在其他 GPU 上。使用这个，想法是将您的批次分成 X 个块（X = GPU 数量）并在每个设备上定义您的模型，每次都将相应的块作为输入并确保重用变量。

如果您正在使用，则此问题中tf.Estimator有一些信息。在这里只需使用两个简单的包装器就可以很容易地进行分布式执行，但我个人无法成功使用它（非常慢并且由于段错误而随机崩溃）。