在TensorFlow中，如何将渐变固定到某个GPU？

Question

我创建了一个模型，在CPU上有一个嵌入层，一个LSTMgpu:0和softmax层gpu:1。在

然后我创建了一个培训计划，如下所示：

optim = tf.train.GradientDescentOptimizer(LEARNING_RATE)
train_op = optim.minimize(final_loss)

训练因内存不足而失败。启用设备放置日志会显示softmax的渐变已放置在gpu:0比gpu:1。如何让TensorFlow正确放置它？在

Answer 1

要将渐变放置在与原始（向前传递）操作相同的设备上，应该将可选参数colocate_gradients_with_ops=True传递给^{}。你的代码看起来像：

optim = tf.train.GradientDescentOptimizer(LEARNING_RATE)
train_op = optim.minimize(final_loss, colocate_gradients_with_ops=True)