首页 文章

如何使用Tensorflow的数据集API将数据移动到多个GPU塔

提问于
浏览 次
9

我们在Tensorflow上运行多GPU作业，并评估从基于队列的模型（使用string_input_producer接口）到新Tensorflow Dataset API的迁移 . 后者似乎提供了一种更简单的方法来同时在Train和Validation之间切换 .

下面的代码片段显示了我们如何做到这一点 .

train_dataset, train_iterator = get_dataset(train_files, batch_size, epochs)
    val_dataset, val_iterator = get_dataset(val_files, batch_size, epochs)


    is_validating = tf.placeholder(dtype=bool, shape=())
    next_batch = tf.cond(is_validating,
               lambda: val_iterator.get_next(),
               lambda: train_iterator.get_next())

    validation_tower = self.num_gpus - 1
    tower_grads = []

    for i in range(self.num_gpus):
        with tf.variable_scope(tf.get_variable_scope(),reuse=(i > 0)):
            with tf.device('/gpu:%d' % i), tf.name_scope('%s_%d' % ('gpu_', i)) as scope:
                if i == validation_tower:
                    images, labels = next_batch
                    # Loss funcs snipped out
                else:
                    images, labels = next_batch
                    # Loss funcs snipped out

get_dataset函数构建数据集，设置映射函数和批处理大小 . 它还构建了一个迭代器，但没有初始化它 . 迭代器的初始化发生在会话开始之前 .

会话运行时提供is_validating布尔值，我们通过feed_dict传递的每个步骤is_validating为True，以使用验证数据集

我的问题是：

假设我有8个gpus，所以我们在7个GPU上进行培训 . 对于这7个GPU中的每一个，迭代器是否从同一点前进，从而为所有7个GPU提供相同的数据？

tensorflow tensorflow-gpu tensorflow-datasets

1 回答

  • 15

    目前有三种主要选择，它们具有不同的可用性和性能权衡：

    • 在Dataset.batch()转换中，创建一个包含所有GPU示例的大批量 . 然后在 Iterator.get_next() 的输出上使用tf.split(..., self.num_gpus)为每个GPU创建子批次 . 这可能是最简单的方法，但它确实将分裂放在关键路径上 .

    • 在 Dataset.batch() 转换中，创建一个适合单个GPU的小批量 . 然后每个GPU调用 Iterator.get_next() 一次以获得多个不同的批次 . （相比之下，在您当前的代码中， next_batch 的相同值被发送到每个GPU，这可能不是您想要发生的 . ）

    • 创建多个迭代器，每个GPU一个 . 在管道的早期使用Dataset.shard()对数据进行分片（例如，如果对数据集进行分片，则在文件列表中） . 请注意，此方法将在主机上消耗更多资源，因此您可能需要调低任何缓冲区大小和/或并行度

    请注意，当前的 tf.data 管道仅在CPU上运行，高效管道的一个重要方面是在上一步仍在运行时将训练输入暂存到GPU . 请参阅TensorFlow CNN benchmarks示例代码，该代码显示如何有效地将数据分段到GPU . 我们目前正在努力直接向 tf.data API添加此支持 .

    回复于

相关问题