Почему несколько IteratorGetNext выполняются последовательно при использовании MirroredStrategy в Tensorflow? - PullRequest
Новая
       12

Почему несколько IteratorGetNext выполняются последовательно при использовании MirroredStrategy в Tensorflow?

0 голосов
/ 29 апреля 2020

Я использую Tensorflow 1.13 для обучения с несколькими GPU на одном узле.

Когда у меня тренировка с 4 GPU в специализированном тренинге l oop, я вижу параллельное выполнение 4 IteratorGetNext на временной шкале:

enter image description here

И мой код для построения конвейера Dataset: 

for i in range(num_gpus):
  dataset = input_dataset.shard(num_gpus, i)
  tower_dataset = total_dataset.shard(num_gpus, i)
  cores = int(multiprocessing.cpu_count() / num_gpus)
  tower_dataset = tower_dataset.interleave(
                        lambda x: tf.data.TFRecordDataset(x, compression_type=compression_type),
                        cycle_length=cores,
                        block_length=cores * 10,
                        num_parallel_calls=cores)
  tower_dataset = tower_dataset.shuffle(...)
  tower_dataset = tower_dataset.repeat(...)
  tower_dataset = tower_dataset.batch(...)
  tower_dataset = tower_dataset.map(parse_function, num_parallel_calls=cores)
  tower_dataset = tower_dataset.prefetch(50)

  iterator = tower_dataset.make_initializable_iterator()
  features.append(iterator.get_next())

dataset_init_op = tf.group(*dataset_init_ops, name="dataset_init_node")
...

Однако, когда я пробую tf.distribute.MirroredStrategy в том же контейнере, я обнаружил, что 4 IteratorGetNext выполнялись последовательно:

enter image description here

И код для этого, вероятно, здесь: 

with strategy.scope():
  with tf.variable_scope(tf.get_variable_scope()) as vscope:
    def model_gpu_fn():
      replica_id = tf.distribute.get_replica_context().replica_id_in_sync_group
      replica_id = tf.cast(replica_id, tf.int64)
      tower_dataset = total_dataset.shard(num_shared, replica_id)
      ...
      tower_dataset = tower_dataset.prefetch(50)
      iterator = tower_dataset.make_initializable_iterator()
      feature = iterator.get_next()
      dataset_init_op = iterator.make_initializer(tower_dataset)
      ...
      return loss
  grads_and_vars = strategy.extended.call_for_each_replica(model_gpu_fn)
train_op = opt._distributed_apply(strategy, grads_and_vars, global_step=global_step)

Эта проблема озадачила меня долгое время. Чем обусловлено серийное исполнение?

...