Question

我在4 GPU机器上使用TF r1.8.0运行训练，并且我尝试使用tf.data和高级TF估算器替换现有的训练代码。我之前的代码主要遵循此处的多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 - 具体来说，我＆＃39; m使用＃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。

将tf.data API与分布式TF

1 个答案: