使用带有分布式 TF 的 tf.data API答案

【问题标题】：using the tf.data API with distributed TF使用带有分布式 TF 的 tf.data API
【发布时间】：2018-05-08 06:35:17
【问题描述】：

我正在使用 TF r1.8.0 在 4 GPU 机器上运行训练，并且我正在尝试使用 tf.data 和高级 TF Estimator 替换我现有的训练代码。我之前的代码主要遵循此处找到的多 GPU CIFAR10 示例代码：https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py

当我用 tf.data 替换现有的输入管道（基于队列）时，我对数据集进行了分片并为每个设备创建了一个迭代器（按照这个问题的答案的建议：How does one move data to multiple GPU towers using Tensorflow's Dataset API - 具体来说，我是使用＃3），一切都很好。

现在，为了利用 tf.contrib.distribute 中的 MirroredStrategy，看来我必须改用 Estimators（除非我弄错了？），但我的问题是：我还需要根据我将使用的 GPU 数量对我的数据集进行分片，或者我是否像在单个设备上一样编写它并信任 Estimator 按 GPU 数量分割每个批次？由于整个训练循环都被抽象掉了，我有点难以理解 Estimator 到底在做什么......

如果这已在某处清楚地记录或以前被问过，我提前道歉！ fwiw，我当前的输入管道如下所示：

def input_fn(tfrecords_dirpath, num_gpus, batch_size, 
             num_epochs, gpu_device, gpu_index):

    tfrecord_filepaths = tf.data.Dataset.list_files('{}/*.tfrecord'.format(tfrecords_dirpath))
    dataset = tf.data.TFRecordDataset(tfrecord_filepaths, num_parallel_reads= int(64 / num_gpus))

    dataset = dataset.shard(num_gpus, gpu_index)

    # use fused operations (shuffle_and_repeat, map_and_batch)
    dataset = dataset.apply(tf.contrib.data.shuffle_and_repeat(10000, num_epochs))
    dataset = dataset.apply(tf.contrib.data.map_and_batch(lambda x: parse_record(x), batch_size))

    # stage batches for processing by loading them pre-emptively on the GPU
    dataset = dataset.apply(tf.contrib.data.prefetch_to_device(gpu_device))

    iterator = dataset.make_one_shot_iterator()
    images_batch, labels_batch = iterator.get_next()

    return images_batch, labels_batch

当我开始训练时，我会在每个 GPU 中复制模型并汇总损失：

# create a separate inference graph in every GPU
gpu_devices = ['/gpu:{}'.format(i) for i in range(num_gpus)]
with tf.variable_scope(tf.get_variable_scope()):
    for i, gpu_device in enumerate(gpu_devices):

        # create a dataset and iterator per GPU
        image_batch, label_batch = input_fn(tfrecords_dirpath, num_gpus, batch_size_per_tower, 
                                            num_epochs, gpu_device, i)
         with tf.device(gpu_device):
            with tf.name_scope('{}_{}'.format('tower', i)) as scope:

                # run inference and compute tower losses
                ...

谢谢！

【问题讨论】：

标签： tensorflow deep-learning multi-gpu tensorflow-estimator

【解决方案1】：

在一台机器上，您不需要分片。 Here 是使用 tf.distribute.MirroredStrategy 和 tf.estimator.train_and_evaluate 的示例。在此设置中，应使用每个 GPU 的批量大小创建数据集对象，并且 TF Estimator 将在每次迭代时在每个 GPU 上运行它。因此，如果每批 GPU 为 B，GPU 数量为 N，则全局批大小将为 N*B。

【讨论】：