Tensorflow 将tf.data API与分布式tf一起使用

我正在4GPU机器上运行使用TFR1.8.0的训练，并试图使用TF.data和高级TF估计器替换现有的训练代码。我以前的代码主要遵循以下多GPU CIFAR10示例代码：

当我用tf.data替换现有的输入管道（基于队列）时，我对数据集进行了切分，并为每个设备创建了一个迭代器（按照这个问题答案的建议：-具体地说，我使用的是#3），一切都很顺利

现在，为了利用tf.contrib.distribute中的MirroredStrategy，看起来我必须切换到使用估计器（除非我弄错了？），但我的问题是：我是否仍然需要根据我将要使用的GPU数量来分割数据集，或者我写它的时候就好像我在一台设备上，并且相信估计器会根据GPU的数量来分割每一批？由于整个训练循环都被抽象掉了，我有点难以理解估计器在幕后到底在做什么

如果这已经被清楚地记录在某个地方或问过，我提前道歉！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
                ...

---

谢谢

在一台机器上，您不需要切分。是使用

tf.distribute.MirroredStrategy

和

tf.estimator.train\u和\u evaluate

时的一个示例。在此设置中，应使用每个GPU批量大小创建dataset对象，TF Estimator将在每次迭代的每个GPU上运行它。因此，如果每批GPU是B，GPU的数量是N，那么全局批大小将是N*B