python - 如何将梯度固定到 TensorFlow 中的某个 GPU？

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

然后我创建了一个训练操作，如下所示:

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

训练失败并出现内存不足错误。启用设备放置日志显示 softmax 的梯度已放置在 gpu:0 而不是 gpu:1 上。如何让 TensorFlow 正确放置它？

最佳答案

要将梯度与原始(前向传递)操作放在同一设备上，您应该将可选参数 colocate_gradients_with_ops=True 传递给 optim.minimize() 。您的代码将如下所示:

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

关于python - 如何将梯度固定到 TensorFlow 中的某个 GPU？，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/36124790/