Tensorflow: MultiGPU Обучение на части графика
Весь код предполагает использование Tensorflow 1.3 и Python 3.x
Мы работаем над алгоритмом GAN, который имеет интересную функцию потерь.
Stage 1 - Compute only the completion/generator loss portion of the network
Iterates over the completion portion of the GAN for X iterations.
Stage 2 - Compute only the discriminator loss portion of the network
Iterates over the discriminator portion for Y iterations (but
don't train on Stage 1)
Stage 3 - Compute the full loss on the network
Iterate over both completion and discriminator for Z iterations
(training on the entire network).
У нас есть этот рабочий сингл GPU. Мы хотим, чтобы он работал на нескольких графических процессорах, так как время тренировок велико.
Мы рассмотрели Tensorflow/models/tutorials/Images/cifar10/cifar10_multi_gpu_train.py, в котором говорится о потере башни, усредняющих башни вместе, вычисляя ваши градиенты на графических процессорах, а затем применяя их к ЦП. Это отличное начало. Однако, поскольку наша потеря более сложна, она все немного усложнила для нас.
Код довольно сложный, но примерно такой же, как https://github.com/timsainb/Tensorflow-MultiGPU-VAE-GAN, (но он не будет работать, потому что он был написан вокруг Tensorflow 0.1, поэтому у него есть некоторые Странности, которые я не получил работать, но это должно дать вам представление о том, что мы делаем)
Когда мы вычисляем градиенты, это выглядит примерно так (псевдокод, чтобы попытаться выделить важные части):
for i in range(num_gpus):
with tf.device('/gpu:%d' % gpus[i]):
with tf.name_scope('Tower_%d' % gpus[i]) as scope:
with tf.variable_scope( "generator" )
generator = build_generator()
with tf.variable_scope( "discriminator" ):
with tf.variable_scope( "real_discriminator" ) :
real_discriminator = build_discriminator(x)
with tf.variable_scope( "fake_discriminator", reuse = True ):
fake_discriminator = build_discriminator(generator)
gen_only_loss, discm_only_loss, full_loss = build_loss( generator,
real_discriminator, fake_discriminator )
tf.get_variable_scope().reuse_variables()
gen_only_grads = gen_only_opt.compute_gradients(gen_only_loss)
tower_gen_only_grads.append(gen_only_grads)
discm_only_train_vars= tf.get_collection(
tf.GraphKeys.TRAINABLE_VARIABLES, "discriminator" )
discm_only_train_vars= discm_only_train_vars+ tf.get_collection(
tf.GraphKeys.TRAINABLE_RESOURCE_VARIABLES, "discriminator" )
discm_only_grads = discm_only_opt.compute_gradients(discm_only_loss,
var_list = discm_only_train_vars)
tower_discm_only_grads.append(discm_only_grads)
full_grads = full_opt.compute_gradients(full_loss)
tower_full_grads.append(full_grads)
# average_gradients is the same code from the cifar10_multi_gpu_train.py.
We haven't changed it. Just iterates over gradients and averages
them...this is part of the problem...
gen_only_grads = average_gradients(tower_gen_only_grads)
gen_only_train = gen_only_opt.apply_gradients(gen_only_grads,
global_step=global_step)
discm_only_grads = average_gradients(tower_discm_only_grads)
discm_only_train = discm_only_opt.apply_gradients(discm_only_grads,
global_step=global_step)
full_grads = average_gradients(tower_full_grads)
full_train = full_opt.apply_gradients(full_grads, global_step=global_step)
Если мы вызываем только "compute_gradients(full_loss)", алгоритм работает правильно на нескольких графических процессорах. Это довольно эквивалентно коду в примере cifar10_multi_gpu_train.py. Сложная часть возникает, когда необходимо ограничить сеть на этапе 1 или 2.
Compute_gradients (full_loss), имеет параметр var_list со значением по умолчанию None, что означает, что он обучает все переменные. Как он узнает, что не нужно тренировать переменные Tower_0 в Tower_1? Я спрашиваю, потому что, когда мы имеем дело с compute_gradients( discm_only_loss, var_list = discm_only_train_vars), мне нужно знать, как собрать правильные переменные, чтобы ограничить обучение этой частью сети. Я обнаружил, что один поток говорит об этом, но обнаружил, что он является неточным / неполным - "заморозить" некоторые переменные / области действия в tenorflow: stop_gradient против передачи переменных для минимизации.
Причина в том, что если вы посмотрите на код в compute_gradients, то в var_list заполняется комбинация обучаемых переменных и переменных обучаемых ресурсов, когда передается None. Так я и ограничил это. Это все работает правильно, если мы не пытаемся разделить несколько графических процессоров.
Вопрос 1: Теперь, когда я разделил сеть по вышкам, я несу ответственность за сбор текущей башни? Мне нужно добавить такую строку?
discm_only_train_vars= tf.get_collection( tf.GraphKeys.TRAINABLE_VARIABLES, "Tower_{}/discriminator".format( i ) )
discm_only_train_vars= discm_only_train_vars + tf.get_collection( tf.GraphKeys.TRAINABLE_RESOURCE_VARIABLES, "Tower_{}/discriminator".format( i ) )
Чтобы обучить правильные переменные для башни (и гарантировать, что я не пропущу обучение этих переменных?)
Вопрос 2: Вероятно, тот же ответ, что и на вопрос 1. Получить "compute_gradients(gen_only_loss)" немного сложнее... в неотвеченной версии gen_only_loss никогда не трогал дискриминатор, поэтому он активировал тензор в графе, который ему нужен, и все было хорошо. Однако в вышедшей версии, когда я вызываю "compute_gradients", он возвращает градиенты для тензоров, которые еще не активированы - поэтому некоторые записи являются [(None, tf.Variable), (None, tf.Variable)]. Это приводит к сбою average_gradients, потому что он не может преобразовать значение None в Tensor. Это заставляет меня думать, что я должен ограничить их.
Все это сбивает с толку то, что пример cifar и мой пример full_loss не заботятся о тренировках на определенных башнях, но я предполагаю, что как только я укажу var_list, любую магию, которую compute_gradients использует, чтобы узнать, на каких переменных тренироваться какие башни исчезают? Нужно ли беспокоиться о получении каких-либо других переменных?
1 ответ
Для вопроса 1, вы несете ответственность за сбор, если вы делите вручную, да.
Для вопроса 2 вы можете захотеть ограничить вызов compute_gradients или отфильтровать результат.