использование API tf.data с распределенным TF

Question

Я выполняю обучение с использованием TF r1.8.0 на машине с 4 графическими процессорами и пытаюсь заменить свой существующий обучающий код с использованием tf.data и высокоуровневого оценщика TF. Мой предыдущий код в основном следовал примеру кода для нескольких графических процессоров CIFAR10, найденному здесь: https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py

Когда я заменил существующий входной конвейер (на основе очереди) на tf.data, я осветил набор данных и создал итератор для каждого устройства (следуя совету ответа на этот вопрос: Как перемещать данные в несколько графических процессоров, использующих Tensorflow Dataset API - в частности, я использую # 3), и все получилось просто отлично.

Теперь, чтобы воспользоваться MirroredStrategy в tf.contrib.distribute, похоже, что мне нужно переключиться на использование Оценщиков (если я не ошибаюсь?), Но мой вопрос таков: мне все еще нужно чтобы разделить мой набор данных на основе количества графических процессоров, которые я буду использовать, или я пишу его так, как будто я работаю на одном устройстве и доверяю Estimator для разделения каждого пакета на количество графических процессоров? Я изо всех сил пытаюсь понять, что на самом деле делает Оценщик под капотом, так как весь цикл обучения отвлечен ...

Если это где-то было четко задокументировано или спрашивалось ранее, заранее прошу прощения! Кстати, мой текущий входной конвейер выглядит так:

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

и когда я начинаю тренироваться, я копирую модель в каждом графическом процессоре и совокупные потери:

# 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
                ...

Спасибо!