tensorflow - 如何使 tf.gradients() 函数在多个 gpu 中工作？

Question

由于梯度会出现“内存不足”错误，如何在多个 GPU 中运行梯度函数？

    with tf.device('gpu:1'):
        trainable_variables = tf.trainable_variables()
        grads = tf.gradients(self.loss, trainable_variables)

    with tf.device('/gpu:2'):
        var1=trainable_variables[60:107]
        grad1=tf.gradients(self.loss, var1)
    with tf.device('/gpu:3'):
        var2=trainable_variables[57:107]
        grad2=tf.gradients(self.loss, var2)
    with tf.device('/gpu:3'):
        optimizer = tf.train.RMSPropOptimizer(
            self.lrn_rate,
            decay=0.9,
            momentum=0.9,
            epsilon=1.0)
        apply_op1 = optimizer.apply_gradients(
            zip(grads, trainable_variables),
            global_step=self.global_step, name='train_step')

上面的代码只在 gpu:1 中运行，因为 gradients\ 或 gradients_1\ 命名空间只在 gpu:1 中存在，这如何在多个 gpu 中工作？

Answer 1

没关系，有两种解决方案： 1.使用这种风格的代码，测试不同范围的内存。2.重构gradient()函数，使用gradient()函数中的with device()，使不同gpu中的变量初始值。