tensorflow - 使用带有分布式 TF 的 tf.data API

Question

我正在 4 GPU 机器上使用 TF r1.8.0 进行训练，并且我正在尝试使用 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
                ...

谢谢!

Answer 1

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