parallel-processing - Tensorflow - 是否可以手动决定图中的哪些张量进入 GPU 和 CPU？

Question

我在 Keras 中创建了一个同时使用卷积层和 LSTM 层的网络。

而且我读到 tensorflow 不能很好地处理 GPU 中的 LSTM 层:http://minimaxir.com/2017/07/cpu-or-gpu/

事实上，我的网络在 GPU 中变慢了(希望这不是进一步的问题)。

我想让tensorflow把所有的卷积运算都丢给GPU，但是把所有的LSTM运算都留在CPU里。这可能吗？

那么，考虑到我可以找到 tensorflow 图并识别它的每个张量，我可以手动定义每个张量执行操作的位置吗？

Answer 1

在执行 this 之后网络，您可以(如果您使用 Tensorflow)尝试以下技术:

with tf.device('/cpu:0'):
    # model layers which should be run on CPU go here