Tensorflow - Is it possible to manually decide which tensors in a graph go to the GPU and CPU?

user_2097240 0
Viewed 222

I've created a network in Keras that uses both convolutional and LSTM layers.

And I've read that tensorflow can't handle well the LSTM layers in GPU: http://minimaxir.com/2017/07/cpu-or-gpu/

In fact, my network got slow in the GPU (hoping its not a further problem).

I want tensorflow to throw all convolutional operations to the GPU, but keep all the LSTM operations in the CPU. Is that possible?

So, considering I can find the tensorflow graph and identify each of its tensors, can I define manually where each tensor performs its operations?

keras tensorflow deep-learning parallel-processing
2 Answers

Following implementation of this network, you could (if you use Tensorflow) try the following technique:

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