Tensorflow：在图的一部分上进行多 GPU 训练

Question

所有代码都假设使用 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。我们想让它在多 GPU 上工作，因为训练时间很长。

我们查看了 Tensorflow/models/tutorials/Images/cifar10/cifar10_multi_gpu_train.py，其中讨论了塔损失，将塔平均在一起，在 GPU 上计算梯度，然后在 CPU 上应用它们。这是一个很好的开始。但是，由于我们的损失更复杂，所以对我们来说一切都变得复杂了一点。

代码相当复杂，但与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)”，该算法可以在多个 GPU 上正常工作。这与 cifar10_multi_gpu_train.py 示例中的代码非常等效。当需要在第 1 阶段或第 2 阶段限制网络时，棘手的部分就出现了。

Compute_gradients(full_loss)，有一个 var_list 参数，默认值为 None，这意味着它训练所有变量。它怎么知道在 Tower_1 中不训练 Tower_0 变量？我问，因为当我们处理 compute_gradients(discm_only_loss, var_list = discm_only_train_vars) 时，我需要知道如何收集正确的变量以将训练限制在网络的该部分。我发现一个线程在谈论这个，但发现它不准确/不完整 - "freeze" some variables/scopes in tensorflow: stop_gradient vs passing variables to minimize。

原因是，如果您查看 compute_gradients 中的代码，var_list 填充的是可训练变量和可训练资源变量的组合，当传入 None 时。所以这也是我对其进行限制的方式。如果我们不尝试跨多个 GPU 拆分，这一切都可以正常工作。

问题 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 用来知道哪些变量的任何魔法训练哪些塔消失了？我需要担心获取任何其他变量吗？

Answer 1

对于问题1，如果你手动拆分，你负责收集，是的。

对于问题 2，您可能希望限制对 compute_gradients 的调用或过滤结果。