TensorFlow на нескольких графических процессорах

Question

В последнее время я пытаюсь научиться использовать Tensorflow на нескольких графических процессорах, прочитав официальный учебник. Однако есть кое-что, что меня смущает. Следующий код является частью официального руководства, которое рассчитывает потери на одном GPU.

def tower_loss(scope, images, labels):

  # Build inference Graph.
  logits = cifar10.inference(images)

  # Build the portion of the Graph calculating the losses. Note that we will
  # assemble the total_loss using a custom function below.
  _ = cifar10.loss(logits, labels)

  # Assemble all of the losses for the current tower only.
  losses = tf.get_collection('losses', scope)

  # Calculate the total loss for the current tower.
  total_loss = tf.add_n(losses, name='total_loss')

  # Attach a scalar summary to all individual losses and the total loss; do the
  # same for the averaged version of the losses.
  for l in losses + [total_loss]:
    # Remove 'tower_[0-9]/' from the name in case this is a multi-GPU training
    # session. This helps the clarity of presentation on tensorboard.
    loss_name = re.sub('%s_[0-9]*/' % cifar10.TOWER_NAME, '', l.op.name)
    tf.summary.scalar(loss_name, l)

  return total_loss

Процесс обучения заключается в следующем.

def train():
with tf.device('/cpu:0'):
    # Create a variable to count the number of train() calls. This equals the
    # number of batches processed * FLAGS.num_gpus.
global_step = tf.get_variable(
    'global_step', [],
    initializer=tf.constant_initializer(0), trainable=False)

# Calculate the learning rate schedule.
num_batches_per_epoch = (cifar10.NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN /
                         FLAGS.batch_size / FLAGS.num_gpus)
decay_steps = int(num_batches_per_epoch * cifar10.NUM_EPOCHS_PER_DECAY)

# Decay the learning rate exponentially based on the number of steps.
lr = tf.train.exponential_decay(cifar10.INITIAL_LEARNING_RATE,
                                global_step,
                                decay_steps,
                                cifar10.LEARNING_RATE_DECAY_FACTOR,
                                staircase=True)

# Create an optimizer that performs gradient descent.
opt = tf.train.GradientDescentOptimizer(lr)

# Get images and labels for CIFAR-10.
images, labels = cifar10.distorted_inputs()
batch_queue = tf.contrib.slim.prefetch_queue.prefetch_queue(
      [images, labels], capacity=2 * FLAGS.num_gpus)
# Calculate the gradients for each model tower.
tower_grads = []
with tf.variable_scope(tf.get_variable_scope()):
  for i in xrange(FLAGS.num_gpus):
    with tf.device('/gpu:%d' % i):
      with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:
        # Dequeues one batch for the GPU
        image_batch, label_batch = batch_queue.dequeue()
        # Calculate the loss for one tower of the CIFAR model. This function
        # constructs the entire CIFAR model but shares the variables across
        # all towers.
        loss = tower_loss(scope, image_batch, label_batch)

        # Reuse variables for the next tower.
        tf.get_variable_scope().reuse_variables()

        # Retain the summaries from the final tower.
        summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)

Тем не менее, я запутался в цикле for о 'for i in xrange(FLAGS.num_gpus)'. Кажется, мне нужно получить новое пакетное изображение из batch_queue и рассчитать каждый градиент. Я думаю, что этот процесс сериализуется, а не параллельно. Если что-то не так с моим собственным пониманием? Кстати, я также могу использовать итератор для подачи изображения в мою модель, а не в очередь?

Спасибо всем!

Answer 1

Это распространенное заблуждение в модели кодирования Tensorflow. То, что вы показываете здесь, это построение вычислительного графа, а не фактическое выполнение.

Блок:

for i in xrange(FLAGS.num_gpus):
    with tf.device('/gpu:%d' % i):
      with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:
        # Dequeues one batch for the GPU
        image_batch, label_batch = batch_queue.dequeue()
        # Calculate the loss for one tower of the CIFAR model. This function
        # constructs the entire CIFAR model but shares the variables across
        # all towers.
        loss = tower_loss(scope, image_batch, label_batch)

переводится как:

For each GPU device (`for i in range..` & `with device...`):
    - build operations needed to dequeue a batch
    - build operations needed to run the batch through the network and compute the loss

Обратите внимание, как через tf.get_variable_scope().reuse_variables() Вы говорите графику, что переменные, используемые для графического графического процессора, должны быть общими для всех (т. е. все графики на нескольких устройствах "повторно" используют одни и те же переменные).

Ничто из этого не запускает сеть один раз (обратите внимание, что нет sess.run()): вы просто даете инструкции о том, как должны передаваться данные.

Затем, когда вы начнете само обучение (я полагаю, что вы пропустили этот фрагмент кода при копировании здесь), каждый графический процессор будет извлекать свою собственную партию и производить потерю для каждой башни. Я предполагаю, что эти потери усредняются где-то в последующем коде, а среднее значение - это потеря, переданная оптимизатору.

Вплоть до того момента, когда потери в башне усредняются вместе, все не зависит от других устройств, поэтому получение пакета и вычисление потерь можно выполнять параллельно. Затем обновление градиентов и параметров выполняется только один раз, переменные обновляются и цикл повторяется.

Таким образом, чтобы ответить на ваш вопрос, нет, расчеты потерь для каждой партии не сериализуются, но, поскольку это синхронные распределенные вычисления, вам нужно собрать все потери со всех графических процессоров, прежде чем продолжить работу по вычислению градиентов и обновлению параметров, так что у вас все еще есть некоторая часть графика, которая не может быть независимой.